from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast
from codegen.scope import Scope, ScopeTree
if TYPE_CHECKING:
from collections.abc import Callable, Iterator
@dataclass(frozen=True)
class NodePath:
parts: tuple[str, ...] = ()
@classmethod
def of(cls, *segments: str) -> NodePath:
return cls(tuple(segments))
@property
def is_root(self) -> bool:
return len(self.parts) == 0
@property
def parent(self) -> NodePath | None:
if self.is_root:
return None
return NodePath(self.parts[: len(self.parts) - 1])
@property
def parents(self) -> Iterator[NodePath]:
for i in range(len(self.parts) - 1, -1, -1):
yield NodePath(self.parts[:i])
@property
def name(self) -> str:
return self.parts[-1] if self.parts else ""
def child(self, segment: str) -> NodePath:
return NodePath(parts=(*self.parts, segment))
def is_relative_to(self, other: NodePath) -> bool:
n = len(other.parts)
return len(self.parts) >= n and self.parts[:n] == other.parts
@staticmethod
def common_ancestor(*paths: NodePath) -> NodePath | None:
if not paths:
return None
common: tuple[str, ...] = paths[0].parts
for p in paths[1:]:
i = 0
limit = min(len(common), len(p.parts))
while i < limit and common[i] == p.parts[i]:
i += 1
common = common[:i]
return NodePath(common)
def id_str(self, sep: str = ".") -> str:
return sep.join(self.parts)
@classmethod
def parse(cls, s: str, sep: str = ".") -> NodePath:
return cls(tuple(s.split(sep))) if s else cls(())
@classmethod
def combine(cls, parent: NodePath | None, s: str) -> NodePath:
suffix = cls.parse(s)
if parent is None:
return suffix
return cls((*parent.parts, *suffix.parts))
def __str__(self) -> str:
return "/".join(self.parts) if self.parts else "<root>"
[docs]
@dataclass
class BuildStore:
scope_tree: ScopeTree = field(default_factory=ScopeTree)
_items: dict[tuple[NodePath, str], list[object]] = field(
default_factory=dict
)
_instances: dict[NodePath, object] = field(default_factory=dict)
def add(
self,
instance_id: NodePath,
op_name: str,
*objects: object,
) -> None:
self._items.setdefault((instance_id, op_name), []).extend(objects)
def register_instance(
self,
instance_id: NodePath,
instance: object,
) -> None:
self._instances[instance_id] = instance
def scope_of(self, instance_id: NodePath) -> Scope:
return self.scope_tree.scope_for(instance_id)
def instance_at(self, instance_id: NodePath) -> object | None:
return self._instances.get(instance_id)
def ancestor_id_of(
self,
instance_id: NodePath,
scope_name: str,
) -> NodePath | None:
for parent in instance_id.parents:
if self.scope_of(parent).name == scope_name:
return parent
return None
def ancestor_of(
self,
instance_id: NodePath,
scope_name: str,
) -> object | None:
ancestor_id = self.ancestor_id_of(instance_id, scope_name)
if ancestor_id is None:
return None
return self._instances.get(ancestor_id)
def common_ancestor_id_of(self, *instance_ids: NodePath) -> NodePath | None:
return NodePath.common_ancestor(*instance_ids)
def children(
self,
parent_id: NodePath,
*,
child_scope: str | None = None,
) -> list[tuple[NodePath, object]]:
out: list[tuple[NodePath, object]] = []
for child_id, instance in self._instances.items():
if child_id.parent != parent_id:
continue
if (
child_scope is not None
and self.scope_of(child_id).name != child_scope
):
continue
out.append((child_id, instance))
return out
def _filter_for_callable[T](
self,
instances: list[T],
filter_fn: Callable[[T], bool] = lambda _: True,
) -> list[T]:
return [instance for instance in instances if filter_fn(instance)]
def _return_first[T](
self,
instances: list[T],
) -> T:
return instances[0]
def outputs_under[T](
self,
instance_id: NodePath,
output_type: type[T] | None = None,
filter_fn: Callable[[T], bool] = lambda _: True,
scope: str | None = None,
) -> list[T]:
instance_callable: Callable[[Any], bool] = (
(lambda obj: isinstance(obj, output_type))
if output_type
else lambda _: True
)
def parsed_callable(obj: Any) -> bool: # noqa
return instance_callable(obj) and filter_fn(obj)
result: list[T] = []
for (stored_id, _), items in self._items.items():
if not stored_id.is_relative_to(instance_id):
continue
if scope is not None and self.scope_of(stored_id).name != scope:
continue
result.extend(cast("T", item) for item in items)
return self._filter_for_callable(result, parsed_callable)
def output_under[T](
self,
instance_id: NodePath,
output_type: type[T] | None = None,
filter_fn: Callable[[T], bool] = lambda _: True,
scope: str | None = None,
) -> T:
return self._return_first(
self.outputs_under(
instance_id=instance_id,
output_type=output_type,
filter_fn=filter_fn,
scope=scope,
)
)
def outputs_under_ancestor[T](
self,
instance_id: NodePath,
ancestor_scope: str,
output_type: type[T] | None = None,
filter_fn: Callable[[T], bool] = lambda _: True,
scope: str | None = None,
) -> list[T]:
ancestor_id = self.ancestor_id_of(instance_id, ancestor_scope)
if ancestor_id is None:
return []
return self.outputs_under(ancestor_id, output_type, filter_fn, scope)
def output_under_ancestor[T](
self,
instance_id: NodePath,
ancestor_scope: str,
output_type: type[T] | None = None,
filter_fn: Callable[[T], bool] = lambda _: True,
scope: str | None = None,
) -> T:
return self._return_first(
self.outputs_under_ancestor(
instance_id,
ancestor_scope,
output_type=output_type,
filter_fn=filter_fn,
scope=scope,
)
)
def entries(self) -> Iterator[tuple[NodePath, str, list[object]]]:
for (instance_id, op_name), items in self._items.items():
yield instance_id, op_name, items