[wip] Rework load_libraryfile and LazyLibrary using overlays

This commit is contained in:
Jan Petykiewicz 2026-04-21 23:17:31 -07:00
commit e108199bcd
7 changed files with 1204 additions and 616 deletions

373
masque/file/gdsii_lazy.py Normal file
View file

@ -0,0 +1,373 @@
"""
Source-backed lazy GDSII reader using the pure-python klamath path.
This module mirrors the lazy Arrow reader's interface closely enough to share
the same overlay and ports-import helpers, while still materializing cells
through the classic `gdsii` decoder.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import IO, Any, cast
from collections import defaultdict
from collections.abc import Iterator, Sequence
import gzip
import io
import logging
import mmap
import pathlib
import klamath
import numpy
from numpy.typing import NDArray
from klamath import records
from . import gdsii
from .utils import is_gzipped
from .gdsii_lazy_core import OverlayLibrary, PortsLibraryView, _pattern_children, write, writefile
from ..error import LibraryError
from ..library import ILibraryView, LibraryView, dangling_mode_t
from ..pattern import Pattern
from ..utils import apply_transforms
logger = logging.getLogger(__name__)
@dataclass
class _SourceHandle:
path: pathlib.Path | None
stream: IO[bytes]
handle: IO[bytes] | None = None
def close(self) -> None:
self.stream.close()
if self.handle is not None and self.handle is not self.stream:
self.handle.close()
self.handle = None
@dataclass(frozen=True)
class _CellScan:
offset: int
children: set[str]
def _open_source_stream(
filename: str | pathlib.Path,
*,
use_mmap: bool,
) -> _SourceHandle:
path = pathlib.Path(filename).expanduser().resolve()
if is_gzipped(path):
if use_mmap:
logger.info('Asked to mmap a gzipped file, reading into memory instead...')
with gzip.open(path, mode='rb') as stream:
data = stream.read()
return _SourceHandle(path=path, stream=io.BytesIO(data))
stream = cast('IO[bytes]', gzip.open(path, mode='rb'))
return _SourceHandle(path=path, stream=stream)
if use_mmap:
handle = path.open(mode='rb', buffering=0)
mapped = cast('IO[bytes]', mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ))
return _SourceHandle(path=path, stream=mapped, handle=handle)
stream = path.open(mode='rb')
return _SourceHandle(path=path, stream=stream)
def _scan_library(
stream: IO[bytes],
) -> tuple[dict[str, Any], list[str], dict[str, _CellScan]]:
library_info = gdsii._read_header(stream)
order: list[str] = []
cells: dict[str, _CellScan] = {}
found_struct = records.BGNSTR.skip_past(stream)
while found_struct:
name = records.STRNAME.skip_and_read(stream).decode('ASCII')
offset = stream.tell()
elements = klamath.library.read_elements(stream)
children = {
element.struct_name.decode('ASCII')
for element in elements
if isinstance(element, klamath.elements.Reference)
}
order.append(name)
cells[name] = _CellScan(offset=offset, children=children)
found_struct = records.BGNSTR.skip_past(stream)
return library_info, order, cells
class GdsLibrarySource(ILibraryView):
"""
Read-only library backed by a seekable GDS stream.
Cells are scanned once up front to discover order and child edges, then
materialized one at a time through the classic `gdsii.read_elements` path.
"""
def __init__(
self,
*,
source: _SourceHandle,
library_info: dict[str, Any],
cell_order: Sequence[str],
cells: dict[str, _CellScan],
) -> None:
self.path = source.path
self.library_info = library_info
self._source = source
self._cell_order = tuple(cell_order)
self._cells = cells
self._cache: dict[str, Pattern] = {}
self._lookups_in_progress: list[str] = []
@classmethod
def from_file(
cls,
filename: str | pathlib.Path,
*,
use_mmap: bool = True,
) -> GdsLibrarySource:
source = _open_source_stream(filename, use_mmap=use_mmap)
source.stream.seek(0)
library_info, cell_order, cells = _scan_library(source.stream)
return cls(source=source, library_info=library_info, cell_order=cell_order, cells=cells)
def __getitem__(self, key: str) -> Pattern:
return self._materialize_pattern(key, persist=True)
def __iter__(self) -> Iterator[str]:
return iter(self._cell_order)
def __len__(self) -> int:
return len(self._cell_order)
def __contains__(self, key: object) -> bool:
return key in self._cells
def source_order(self) -> tuple[str, ...]:
return self._cell_order
def materialize_many(
self,
names: Sequence[str],
*,
persist: bool = True,
) -> LibraryView:
mats = {
name: self._materialize_pattern(name, persist=persist)
for name in dict.fromkeys(names)
}
return LibraryView(mats)
def _materialize_pattern(self, name: str, *, persist: bool) -> Pattern:
if name in self._cache:
return self._cache[name]
if name not in self._cells:
raise KeyError(name)
if name in self._lookups_in_progress:
chain = ' -> '.join(self._lookups_in_progress + [name])
raise LibraryError(
f'Detected circular reference or recursive lookup of "{name}".\n'
f'Lookup chain: {chain}\n'
'This may be caused by an invalid (cyclical) reference, or buggy code.\n'
'If you are lazy-loading a file, try a non-lazy load and check for reference cycles.'
)
self._lookups_in_progress.append(name)
try:
self._source.stream.seek(self._cells[name].offset)
pat = gdsii.read_elements(self._source.stream, raw_mode=True)
finally:
self._lookups_in_progress.pop()
if persist:
self._cache[name] = pat
return pat
def _raw_children(self, name: str) -> set[str]:
return set(self._cells[name].children)
def child_graph(
self,
dangling: dangling_mode_t = 'error',
) -> dict[str, set[str]]:
graph: dict[str, set[str]] = {}
for name in self._cell_order:
if name in self._cache:
graph[name] = _pattern_children(self._cache[name])
else:
graph[name] = self._raw_children(name)
existing = set(graph)
dangling_refs = set().union(*(children - existing for children in graph.values()))
if dangling == 'error':
if dangling_refs:
raise self._dangling_refs_error(cast('set[str]', dangling_refs), 'building child graph')
return graph
if dangling == 'ignore':
return {name: {child for child in children if child in existing} for name, children in graph.items()}
for child in dangling_refs:
graph.setdefault(cast('str', child), set())
return graph
def parent_graph(
self,
dangling: dangling_mode_t = 'error',
) -> dict[str, set[str]]:
child_graph = self.child_graph(dangling='include' if dangling == 'include' else 'ignore')
existing = set(self.keys())
igraph: dict[str, set[str]] = {name: set() for name in child_graph}
for parent, children in child_graph.items():
for child in children:
if child in existing or dangling == 'include':
igraph.setdefault(child, set()).add(parent)
if dangling == 'error':
raw = self.child_graph(dangling='include')
dangling_refs = set().union(*(children - existing for children in raw.values()))
if dangling_refs:
raise self._dangling_refs_error(cast('set[str]', dangling_refs), 'building parent graph')
return igraph
def subtree(
self,
tops: str | Sequence[str],
) -> ILibraryView:
if isinstance(tops, str):
tops = (tops,)
keep = cast('set[str]', self.referenced_patterns(tops) - {None})
keep |= set(tops)
return self.materialize_many(tuple(keep), persist=True)
def tops(self) -> list[str]:
graph = self.child_graph(dangling='ignore')
names = set(graph)
not_toplevel: set[str] = set()
for children in graph.values():
not_toplevel |= children
return list(names - not_toplevel)
def with_ports_from_data(
self,
*,
layers: Sequence[tuple[int, int] | int],
max_depth: int = 0,
skip_subcells: bool = True,
) -> PortsLibraryView:
return PortsLibraryView(
self,
layers=layers,
max_depth=max_depth,
skip_subcells=skip_subcells,
)
def find_refs_local(
self,
name: str,
parent_graph: dict[str, set[str]] | None = None,
dangling: dangling_mode_t = 'error',
) -> dict[str, list[NDArray[numpy.float64]]]:
instances: dict[str, list[NDArray[numpy.float64]]] = defaultdict(list)
if parent_graph is None:
graph_mode = 'ignore' if dangling == 'ignore' else 'include'
parent_graph = self.parent_graph(dangling=graph_mode)
if name not in self:
if name not in parent_graph:
return instances
if dangling == 'error':
raise self._dangling_refs_error({name}, f'finding local refs for {name!r}')
if dangling == 'ignore':
return instances
for parent in parent_graph.get(name, set()):
if parent in self._cache:
for ref in self._cache[parent].refs.get(name, []):
instances[parent].append(ref.as_transforms())
continue
pat = self._materialize_pattern(parent, persist=False)
for ref in pat.refs.get(name, []):
instances[parent].append(ref.as_transforms())
return instances
def find_refs_global(
self,
name: str,
order: list[str] | None = None,
parent_graph: dict[str, set[str]] | None = None,
dangling: dangling_mode_t = 'error',
) -> dict[tuple[str, ...], NDArray[numpy.float64]]:
graph_mode = 'ignore' if dangling == 'ignore' else 'include'
if order is None:
order = self.child_order(dangling=graph_mode)
if parent_graph is None:
parent_graph = self.parent_graph(dangling=graph_mode)
if name not in self:
if name not in parent_graph:
return {}
if dangling == 'error':
raise self._dangling_refs_error({name}, f'finding global refs for {name!r}')
if dangling == 'ignore':
return {}
self_keys = set(self.keys())
transforms: dict[str, list[tuple[tuple[str, ...], NDArray[numpy.float64]]]]
transforms = defaultdict(list)
for parent, vals in self.find_refs_local(name, parent_graph=parent_graph, dangling=dangling).items():
transforms[parent] = [((name,), numpy.concatenate(vals))]
for next_name in order:
if next_name not in transforms:
continue
if not parent_graph.get(next_name, set()) & self_keys:
continue
outers = self.find_refs_local(next_name, parent_graph=parent_graph, dangling=dangling)
inners = transforms.pop(next_name)
for parent, outer in outers.items():
outer_tf = numpy.concatenate(outer)
for path, inner in inners:
combined = apply_transforms(outer_tf, inner)
transforms[parent].append(((next_name,) + path, combined))
result = {}
for parent, targets in transforms.items():
for path, instances in targets:
result[(parent,) + path] = instances
return result
def close(self) -> None:
self._source.close()
def __enter__(self) -> GdsLibrarySource:
return self
def __exit__(self, *_args: object) -> None:
self.close()
def read(
stream: IO[bytes],
) -> tuple[GdsLibrarySource, dict[str, Any]]:
source = _SourceHandle(path=None, stream=stream)
stream.seek(0)
library_info, cell_order, cells = _scan_library(stream)
lib = GdsLibrarySource(source=source, library_info=library_info, cell_order=cell_order, cells=cells)
return lib, library_info
def readfile(
filename: str | pathlib.Path,
*,
use_mmap: bool = True,
) -> tuple[GdsLibrarySource, dict[str, Any]]:
lib = GdsLibrarySource.from_file(filename, use_mmap=use_mmap)
return lib, lib.library_info