masque/masque/library/overlay.py

432 lines
16 KiB
Python

"""Overlay and ports-importing library views."""
from __future__ import annotations
from collections import defaultdict
from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal, Self, cast
import copy
from ..error import LibraryError
from ..pattern import Pattern, map_targets
from .base import ILibrary, ILibraryView
from .capabilities import IBorrowing, IMaterializable
from .utils import INameView, dangling_mode_t, _plan_source_names, _source_rename_map, _validate_dangling_mode
from .mapping import LibraryView
if TYPE_CHECKING:
from collections.abc import Callable, Iterator, Mapping, Sequence
import numpy
from numpy.typing import NDArray
from ..ports import Port
from ..utils import layer_t
@dataclass
class _SourceLayer:
""" One imported source layer tracked by an `OverlayLibrary`. """
library: ILibraryView
source_target_map: dict[str, str]
child_graph: dict[str, set[str]]
@dataclass(frozen=True)
class _SourceEntry:
""" Reference to a single visible source-backed cell in an overlay. """
layer_index: int
source_name: str
def _materialize_detached_pattern(view: ILibraryView, name: str) -> Pattern:
if isinstance(view, IMaterializable):
return view.materialize(name, persist=False).deepcopy()
return view[name].deepcopy()
class PortsLibraryView(ILibraryView, IMaterializable, IBorrowing):
"""
Read-only view which imports or applies ports on first materialization.
The wrapped source remains untouched; this view owns a separate processed
cache so direct-copy workflows can continue to use the raw source view.
The view borrows its source: callers must keep the source open for the
lifetime of the view and close the source themselves.
Graph queries and source ordering are delegated to the wrapped source,
while `source_cell()` exposes unchanged layout provenance and `__getitem__`
and `materialize_many()` return port-imported patterns.
"""
def __init__(
self,
source: ILibraryView,
*,
layers: Sequence[layer_t] = (),
max_depth: int = 0,
skip_subcells: bool = True,
ports: Mapping[str, Mapping[str, Port]] | None = None,
replace: bool = False,
) -> None:
self._source = source
self._layers = tuple(layers)
self._max_depth = max_depth
self._skip_subcells = skip_subcells
self._ports = {
name: copy.deepcopy(dict(cell_ports))
for name, cell_ports in (ports or {}).items()
}
self._replace = replace
self._cache: dict[str, Pattern] = {}
self._lookups_in_progress: list[str] = []
def __getitem__(self, key: str) -> Pattern:
return self.materialize(key, persist=True)
def __iter__(self) -> Iterator[str]:
return iter(self._source)
def __len__(self) -> int:
return len(self._source)
def __contains__(self, key: object) -> bool:
return key in self._source
def materialize(self, name: str, *, persist: bool = True) -> Pattern:
from ..utils.ports2data import data_to_ports # noqa: PLC0415
if name in self._cache:
return self._cache[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.'
)
self._lookups_in_progress.append(name)
try:
pat = _materialize_detached_pattern(self._source, name)
if self._layers:
pat = data_to_ports(
layers=self._layers,
library=self,
pattern=pat,
name=name,
max_depth=self._max_depth,
skip_subcells=self._skip_subcells,
)
if name in self._ports:
ports = copy.deepcopy(self._ports[name])
if self._replace:
pat.ports = ports
else:
pat.ports.update(ports)
finally:
self._lookups_in_progress.pop()
if persist:
self._cache[name] = pat
return pat
def source_order(self) -> tuple[str, ...]:
return self._source.source_order()
def borrowed_sources(self) -> tuple[ILibraryView, ...]:
return (self._source,)
def source_cell(self, name: str) -> tuple[ILibraryView, str] | None:
if name not in self._source or name in self._cache:
return None
return self._source, name
def child_graph(
self,
dangling: dangling_mode_t = 'error',
) -> dict[str, set[str]]:
_validate_dangling_mode(dangling)
return self._source.child_graph(dangling=dangling)
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]]]:
_validate_dangling_mode(dangling)
finder = getattr(self._source, 'find_refs_local', None)
if callable(finder):
return cast('dict[str, list[NDArray[numpy.float64]]]', finder(name, parent_graph=parent_graph, dangling=dangling))
return super().find_refs_local(name, parent_graph=parent_graph, dangling=dangling)
class OverlayLibrary(ILibrary, IMaterializable, IBorrowing):
"""
Mutable overlay over one or more source libraries.
Source-backed cells remain lazy until accessed through `__getitem__`, at
which point that visible cell is promoted into an overlay-owned materialized
`Pattern`.
Source libraries must remain open and must not be mutated after they are
added. The overlay borrows each source and snapshots its names, hierarchy,
and initial visible-name mapping while retaining the source itself for lazy
pattern materialization.
"""
def __init__(self) -> None:
self._layers: list[_SourceLayer] = []
self._entries: dict[str, Pattern | _SourceEntry] = {}
self._order: list[str] = []
self._target_remap: dict[str, str] = {}
def __iter__(self) -> Iterator[str]:
return (name for name in self._order if name in self._entries)
def __len__(self) -> int:
return len(self._entries)
def __contains__(self, key: object) -> bool:
return key in self._entries
def __getitem__(self, key: str) -> Pattern:
return self.materialize(key, persist=True)
def __setitem__(
self,
key: str,
value: Pattern | Callable[[], Pattern],
) -> None:
if key in self._entries:
raise LibraryError(f'"{key}" already exists in the library. Overwriting is not allowed!')
pattern = value() if callable(value) else value
self._entries[key] = pattern
if key not in self._order:
self._order.append(key)
def __delitem__(self, key: str) -> None:
if key not in self._entries:
raise KeyError(key)
del self._entries[key]
def _merge(self, key_self: str, other: Mapping[str, Pattern], key_other: str) -> None:
self[key_self] = copy.deepcopy(other[key_other])
def add_source(
self,
source: Mapping[str, Pattern] | ILibraryView,
*,
rename_theirs: Callable[[INameView, str], str] | None = None,
rename_when: Literal['conflict', 'always'] = 'conflict',
) -> dict[str, str]:
"""
Add a source-backed library layer.
The source must remain open, and its names, hierarchy, and pattern
contents must remain unchanged for the lifetime of this overlay.
Args:
rename_theirs: Function used to choose visible names for imported
source cells. Its `INameView` argument contains existing and
previously reserved names, but does not support pattern lookup.
rename_when: If `'conflict'`, only conflicting names are renamed.
If `'always'`, every imported source name is passed through
`rename_theirs`.
"""
view = source if isinstance(source, ILibraryView) else LibraryView(source)
source_order = list(view.source_order())
child_graph = view.child_graph(dangling='include')
source_to_visible = _plan_source_names(
self,
source_order,
rename_theirs = rename_theirs,
rename_when = rename_when,
)
layer = _SourceLayer(
library=view,
source_target_map=dict(source_to_visible),
child_graph=child_graph,
)
layer_index = len(self._layers)
self._layers.append(layer)
for source_name, visible_name in source_to_visible.items():
self._entries[visible_name] = _SourceEntry(layer_index=layer_index, source_name=source_name)
if visible_name not in self._order:
self._order.append(visible_name)
return _source_rename_map(source_to_visible)
def rename(
self,
old_name: str,
new_name: str,
move_references: bool = False,
) -> OverlayLibrary:
if old_name not in self._entries:
raise LibraryError(f'"{old_name}" does not exist in the library.')
if old_name == new_name:
return self
if new_name in self._entries:
raise LibraryError(f'"{new_name}" already exists in the library.')
entry = self._entries.pop(old_name)
self._entries[new_name] = entry
idx = self._order.index(old_name)
self._order[idx] = new_name
if move_references:
self.move_references(old_name, new_name)
return self
def _resolve_target(self, target: str) -> str:
seen: set[str] = set()
current = target
while current in self._target_remap:
if current in seen:
raise LibraryError(f'Cycle encountered while resolving target remap for {target!r}')
seen.add(current)
current = self._target_remap[current]
return current
def _set_target_remap(self, old_target: str, new_target: str) -> None:
resolved_new = self._resolve_target(new_target)
if resolved_new == old_target:
raise LibraryError(f'Ref target remap would create a cycle: {old_target!r} -> {new_target!r}')
self._target_remap[old_target] = resolved_new
for key in list(self._target_remap):
self._target_remap[key] = self._resolve_target(self._target_remap[key])
def move_references(self, old_target: str, new_target: str) -> OverlayLibrary:
if old_target == new_target:
return self
self._set_target_remap(old_target, new_target)
for entry in list(self._entries.values()):
if isinstance(entry, Pattern) and old_target in entry.refs:
entry.refs[new_target].extend(entry.refs[old_target])
del entry.refs[old_target]
return self
def _effective_target(self, layer: _SourceLayer, target: str) -> str:
visible = layer.source_target_map.get(target, target)
return self._resolve_target(visible)
def materialize(self, name: str, *, persist: bool = True) -> Pattern:
if name not in self._entries:
raise KeyError(name)
entry = self._entries[name]
if isinstance(entry, Pattern):
return entry
layer = self._layers[entry.layer_index]
source_pat = _materialize_detached_pattern(layer.library, entry.source_name)
def remap(target: str | None) -> str | None:
return None if target is None else self._effective_target(layer, target)
if source_pat.refs:
source_pat.refs = map_targets(source_pat.refs, remap)
pat = source_pat
if persist:
self._entries[name] = pat
return pat
def child_graph(
self,
dangling: dangling_mode_t = 'error',
) -> dict[str, set[str]]:
_validate_dangling_mode(dangling)
graph: dict[str, set[str]] = {}
for name in self._order:
if name not in self._entries:
continue
entry = self._entries[name]
if isinstance(entry, Pattern):
graph[name] = {child for child, refs in entry.refs.items() if child is not None and refs}
continue
layer = self._layers[entry.layer_index]
children = {self._effective_target(layer, child) for child in layer.child_graph.get(entry.source_name, set())}
graph[name] = children
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 subtree(
self,
tops: str | Sequence[str],
) -> Self:
if isinstance(tops, str):
tops = (tops,)
graph = self.child_graph(dangling='include')
keep = self._referenced_patterns_from_graph(graph, tops=tops)
keep &= set(self)
keep |= set(tops)
new = type(self)()
new._layers = [
_SourceLayer(
library=layer.library,
source_target_map=dict(layer.source_target_map),
child_graph={name: set(children) for name, children in layer.child_graph.items()},
)
for layer in self._layers
]
new._order = [name for name in self._order if name in keep and name in self._entries]
new._entries = {name: self._entries[name] for name in new._order}
new._target_remap = dict(self._target_remap)
return new
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]]]:
_validate_dangling_mode(dangling)
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()):
pat = self.materialize(parent, persist=False)
for ref in pat.refs.get(name, []):
instances[parent].append(ref.as_transforms())
return instances
def source_order(self) -> tuple[str, ...]:
return tuple(name for name in self._order if name in self._entries)
def borrowed_sources(self) -> tuple[ILibraryView, ...]:
return tuple(layer.library for layer in self._layers)
def source_cell(self, name: str) -> tuple[ILibraryView, str] | None:
entry = self._entries.get(name)
if not isinstance(entry, _SourceEntry):
return None
layer = self._layers[entry.layer_index]
children = layer.child_graph.get(entry.source_name, set())
if any(self._effective_target(layer, child) != child for child in children):
return None
return layer.library, entry.source_name