Source code for jitx.query
"""
JITX Transformation Query Layer
================================
A pluggable layer on top of :py:func:`~jitx.inspect.visit` for queries that
need conversion. Where :py:func:`~jitx.inspect.visit` finds objects already
present in the tree of the requested type, :py:class:`TransformQuery` walks
the tree and applies registered transformer functions to convert source
objects into the requested target type.
A transformer is a function ``(Trace, S) -> Iterable[(Trace, D)]`` registered
with :py:meth:`TransformQuery.transformer`. The engine maintains an adjacency
map ``src -> [Transformer]`` updated on each registration, and at query time
recursively walks forward from each found instance, halting when the recursion
produces an instance that matches one of the query's targets.
A module-level :py:data:`default` instance backs the convenience functions
:py:func:`query` and :py:func:`transformer`. Built-in transformers (such as
``Pad -> Copper``, ``Via -> Copper``, ``Via -> Cutout``) live in the source
files where their source types are defined and self-register on the default at
import time. Independent :py:class:`TransformQuery` instances can be
constructed for isolated registries.
Note that transformers do **not** apply ``trace.transform`` to their output:
the produced object is in the source's local coordinate frame, and the caller
composes ``trace.transform`` with the object's geometry if/when needed. This
keeps the layer pluggable for any user-defined target type, even ones for
which world-coordinate composition is not meaningful, and propagates a
``None`` transform (for example when ``transform=None`` was passed) faithfully.
A transformer _may_ modify the ``trace.transform`` if it's reasonable to update
the returned object's frame of reference, but it will not be applied to the
output object (if applicable), and it's up to the user of the interface to apply
it if desired.
>>> # Walk a design and gather every Copper on layer 3, regardless of whether
>>> # it lives as an explicit Copper, a Pour, an OverlappableCopper feature,
>>> # or a per-layer Pad shape.
>>> for trace, c in query(design, Copper, filter=lambda c: c.layer == 3):
... ...
"""
from __future__ import annotations
from collections.abc import Callable, Generator, Iterable
from contextlib import contextmanager
from dataclasses import dataclass
from types import GenericAlias, UnionType
from typing import Any, get_args, get_origin, overload
from .inspect import Trace, visit
from .transform import IDENTITY, Transform
type TransformerFn[S, D] = Callable[[Trace, S], Iterable[tuple[Trace, D]]]
[docs]
@dataclass(frozen=True)
class Transformer:
"""A registered transformer from one type to another. The function takes a
:py:class:`~jitx.inspect.Trace` and a source instance and yields
``(Trace, target_instance)`` pairs."""
src: type
dst: type
fn: TransformerFn
def _normalize_targets(
target: type | tuple[type, ...] | UnionType | GenericAlias,
) -> tuple[type, ...]:
"""Normalize ``target`` to a tuple of concrete types. Mirrors
:py:func:`~jitx.inspect.visit`'s normalization: accepts a single type, a
tuple of types, or a :py:class:`types.UnionType`. Generic aliases are
unwrapped to their origin type."""
if isinstance(target, type):
return (target,)
if isinstance(target, UnionType):
targets: tuple[Any, ...] = get_args(target)
elif isinstance(target, tuple):
targets = target
elif get_origin(target) is not None:
origin = get_origin(target)
if isinstance(origin, type):
return (origin,)
raise ValueError(f"Unsupported query target: {target}")
else:
raise ValueError(f"Unsupported query target: {target}")
normalized = tuple(get_origin(t) or t for t in targets)
for t in normalized:
if not isinstance(t, type):
raise ValueError(f"Unsupported query target: {t}")
return normalized
[docs]
class TransformQuery:
"""Registry and engine for transformation queries. Construct the object, add
transformers using the :py:decorator:`~TransformQuery.transformer` and query
a design or object using :py:meth:`~TransformQuery.query`.
Args:
base: If provided, its transformations will be copied and reused in this
registry. It's a one-time copy, and will not reflect subsequent
changes to the base transform query.
"""
def __init__(self, base: TransformQuery | None = None):
self._transformers: list[Transformer] = list(base._transformers) if base else []
self._forward: dict[type, list[Transformer]] = (
dict(base._forward) if base else {}
)
[docs]
def register(self, src: type, dst: type, fn: TransformerFn) -> None:
"""Register ``fn`` as a transformer from ``src`` to ``dst``."""
t = Transformer(src, dst, fn)
self._transformers.append(t)
self._forward.setdefault(src, []).append(t)
[docs]
def transformer[S, D](self, src: type[S], dst: type[D]):
"""Decorator form of :py:meth:`register`.
>>> @engine.transformer(Pad, Copper)
... def pad_to_copper(trace, pad):
... for layer, shape in _explicit_pad_shapes(pad):
... yield trace, Copper(shape, layer)
"""
def decorator(fn: TransformerFn[S, D]) -> TransformerFn[S, D]:
self.register(src, dst, fn)
return fn
return decorator
def __walk(
self,
trace: Trace,
obj: Any,
targets: tuple[type, ...],
opaque: tuple[type, ...] | None,
visited: frozenset[Transformer],
) -> Iterable[tuple[Trace, Any]]:
if isinstance(obj, targets):
yield trace, obj
return
if opaque and isinstance(obj, opaque):
return
for src_type, transformers in self._forward.items():
if isinstance(obj, src_type):
for t in transformers:
if t in visited:
continue
for next_trace, next_obj in t.fn(trace, obj):
yield from self.__walk(
next_trace, next_obj, targets, opaque, visited | {t}
)
@overload
def query[T: UnionType](
self,
root,
target: T,
/,
*,
through: tuple[type, ...] | UnionType | None = None,
transform: Transform | None = IDENTITY,
opaque: type | tuple[type, ...] | UnionType | None = None,
refs: bool = False,
filter: Callable[[T], bool] | None = None,
) -> Generator[tuple[Trace, T], None, None]: ...
@overload
def query[T](
self,
root,
target: type[T] | tuple[type[T], ...],
/,
*,
through: tuple[type, ...] | UnionType | None = None,
transform: Transform | None = IDENTITY,
opaque: type | tuple[type, ...] | UnionType | None = None,
refs: bool = False,
filter: Callable[[T], bool] | None = None,
) -> Generator[tuple[Trace, T], None, None]: ...
[docs]
def query(
self,
root,
target: type | tuple[type, ...] | UnionType,
/,
*,
through: tuple[type, ...] | UnionType | None = None,
transform: Transform | None = IDENTITY,
opaque: type | tuple[type, ...] | UnionType | None = None,
refs: bool = False,
filter: Callable[[Any], bool] | None = None,
):
"""Walk ``root`` and yield ``(Trace, target)`` for every instance of
``target`` produced by the registered transformer graph.
``target`` may be a single type, a tuple of types, or a
:py:class:`types.UnionType` (e.g. ``Copper | Pour``), mirroring
:py:func:`~jitx.inspect.visit`. Subclasses count as matches.
For each instance found by :py:func:`~jitx.inspect.visit`, the engine
forward-walks registered transformers and yields the first instance
of any target produced. This means direct-target instances yield as
identity (``Pour`` for ``target=Copper``), source-is-target chains
are collapsed (``Pad`` in ``Pad | Copper`` yields as ``Pad``, the
``Pad -> Copper`` chain is bypassed), and longer chains truncate at
the first intermediate that matches a target (with
``Object -> Pour -> Copper`` registered, ``query(design, Pour |
Copper)`` yields the ``Pour`` and never reaches ``Copper``).
Arguments mirror :py:func:`~jitx.inspect.visit`. ``filter`` is applied
to the produced ``target`` instances, not to the source instances.
``trace.transform`` may be ``None`` (when ``transform=None`` was
passed, or when an intermediate element along the path has no
transform); transformers pass it through verbatim.
"""
targets = _normalize_targets(target)
if opaque:
opaque_targets = _normalize_targets(opaque)
else:
opaque_targets = None
# Visit candidates: every registered source type plus the targets
# themselves (so direct-target instances are findable). Some sources
# may dead-end without reaching a target -- the walker silently
# produces no yields for those, a benign over-approximation.
source_types = tuple(set(self._forward.keys()) | set(targets))
if not source_types:
return
@contextmanager
def frame():
from jitx.design import DesignContext
from jitx.substrate import SubstrateContext
from jitx._instantiation import instantiation
with instantiation.require():
# transformers may rely on these to get layer information and such.
with instantiation.frame() as f:
f.disposable = True
with DesignContext(root):
with SubstrateContext(root.substrate):
yield
with frame():
for trace, src_obj in visit(
root,
source_types,
through,
transform=transform,
opaque=opaque_targets,
refs=refs,
):
for out_trace, out_obj in self.__walk(
trace, src_obj, targets, opaque_targets, frozenset()
):
if filter is None or filter(out_obj):
yield out_trace, out_obj
default = TransformQuery()
"""The default :py:class:`TransformQuery` instance used by the module-level
:py:func:`query` and :py:func:`transformer` convenience helpers, and against
which the built-in transformers are registered."""
[docs]
def transformer[S, D](src: type[S], dst: type[D]):
"""Convenience: register on the :py:data:`default` engine.
>>> @transformer(Pad, Copper)
... def pad_to_copper(trace, pad):
... ...
"""
return default.transformer(src, dst)
@overload
def query[T](
root,
target: type[T],
/,
*,
through: tuple[type, ...] | UnionType | None = None,
transform: Transform | None = IDENTITY,
opaque: type | tuple[type, ...] | UnionType | None = None,
refs: bool = False,
filter: Callable[[T], bool] | None = None,
) -> Generator[tuple[Trace, T], None, None]: ...
@overload
def query[T](
root,
target: tuple[type[T], ...],
/,
*,
through: tuple[type, ...] | UnionType | None = None,
transform: Transform | None = IDENTITY,
opaque: type | tuple[type, ...] | UnionType | None = None,
refs: bool = False,
filter: Callable[[T], bool] | None = None,
) -> Generator[tuple[Trace, T], None, None]: ...
[docs]
def query(
root,
target: type | tuple[type, ...] | UnionType,
/,
*,
through: tuple[type, ...] | UnionType | None = None,
transform: Transform | None = IDENTITY,
opaque: type | tuple[type, ...] | UnionType | None = None,
refs: bool = False,
filter: Callable[[Any], bool] | None = None,
):
"""Convenience: query against the :py:data:`default` engine. See
:py:meth:`TransformQuery.query`."""
return default.query(
root,
target,
through=through,
transform=transform,
opaque=opaque,
refs=refs,
filter=filter,
)