class ActorWrapper:
"""Wrapper class for managing actor instances."""
def __init__(self, *args: Any, **kwargs: Any) -> None:
config = parse_root_config(*args, **kwargs)
w = ActorWrapper.from_config(config)
self.actor = w.actor
self.init_ctx = w.init_ctx
@property
def vertex_config(self) -> VertexConfig:
return self.init_ctx.vertex_config
@property
def edge_config(self) -> EdgeConfig:
return self.init_ctx.edge_config
@property
def infer_edges(self) -> bool:
return self.init_ctx.infer_edges
@property
def infer_edge_only(self) -> set[EdgeId]:
return self.init_ctx.infer_edge_only
@property
def infer_edge_except(self) -> set[EdgeId]:
return self.init_ctx.infer_edge_except
@property
def target_db_flavor(self) -> DBType | None:
return self.init_ctx.target_db_flavor
def _inverse_pairs(self) -> dict[str, str]:
"""``{relation: declared inverse}`` of this resource's edge config, built once.
Keyed by the config object so a wrapper re-initialized against another
schema does not mirror through a stale table.
"""
edge_config = self.init_ctx.edge_config
cached = self.__dict__.get("_inverse_pairs_cache")
if cached is None or cached[0] is not edge_config:
cached = (edge_config, inverse_map(edge_config.inverses))
self.__dict__["_inverse_pairs_cache"] = cached
return cached[1]
def init_transforms(self, init_ctx: ActorInitContext) -> None:
self.init_ctx = init_ctx
self.actor.init_transforms(init_ctx)
def finish_init(self, init_ctx: ActorInitContext) -> None:
self.init_ctx = init_ctx
self.actor.init_transforms(init_ctx)
self.actor.finish_init(init_ctx)
# What inference may write: the edges declared, and those the steps
# registered while initialising -- not ones a record registers mid-cast.
self.__dict__["_inferable_edge_ids"] = frozenset(
edge_id for edge_id, _ in init_ctx.edge_config.items()
)
def count(self) -> int:
return self.actor.count()
@classmethod
def from_config(cls, config: ActorConfig) -> ActorWrapper:
if isinstance(config, VertexActorConfig):
actor = VertexActor.from_config(config)
elif isinstance(config, TransformActorConfig):
actor = TransformActor.from_config(config)
elif isinstance(config, EdgeActorConfig):
actor = EdgeActor.from_config(config)
elif isinstance(config, DescendActorConfig):
actor = DescendActor.from_config(config)
elif isinstance(config, VertexRouterActorConfig):
actor = VertexRouterActor.from_config(config)
else:
raise ValueError(
f"Expected VertexActorConfig, TransformActorConfig, EdgeActorConfig, "
f"DescendActorConfig, or VertexRouterActorConfig, got {type(config)}"
)
wrapper = cls.__new__(cls)
wrapper.actor = actor
wrapper.init_ctx = ActorInitContext(
vertex_config=VertexConfig(vertices=[]),
edge_config=EdgeConfig(),
transforms={},
allowed_vertex_names=None,
infer_edges=True,
infer_edge_only=set(),
infer_edge_except=set(),
)
return wrapper
@classmethod
def _from_step(cls, step: dict[str, Any]) -> ActorWrapper:
config = validate_actor_step(normalize_actor_step(step))
return cls.from_config(config)
def __call__(
self,
ctx: ExtractionContext,
lindex: LocationIndex | None = None,
*nargs: Any,
**kwargs: Any,
) -> ExtractionContext:
if lindex is None:
lindex = LocationIndex()
ctx = self.actor(ctx, lindex, *nargs, **kwargs)
return ctx
def assemble(
self, ctx: ExtractionContext | AssemblyContext | ActionContext
) -> defaultdict[GraphEntity, list]:
if isinstance(ctx, AssemblyContext):
assembly_ctx = ctx
else:
assembly_ctx = AssemblyContext.from_extraction(ctx)
# Synthetic identities must exist before edges are assembled and before
# docs are deduplicated on their identity fields: a hash/funnel vertex
# keys on ``id``, so an empty ``id`` gives edges no endpoint key and
# fuse_doc_basis no basis (it would fold the batch into one doc).
ensure_assigned_uuids_in_acc_vertex(assembly_ctx.acc_vertex, self.vertex_config)
ensure_digest_identities_in_acc_vertex(
assembly_ctx.acc_vertex, self.vertex_config
)
assemble_edges(
ctx=assembly_ctx,
vertex_config=self.vertex_config,
edge_config=self.edge_config,
infer_edges=self.infer_edges,
infer_edge_only=self.infer_edge_only,
infer_edge_except=self.infer_edge_except,
target_db_flavor=self.target_db_flavor,
edge_derivation=self.init_ctx.edge_derivation,
inverse_pairs=self._inverse_pairs(),
inferable=self.__dict__.get("_inferable_edge_ids"),
)
for vertex_name, dd in assembly_ctx.acc_vertex.items():
for vertex_list in dd.values():
# Lookup-only observations locate existing vertices for edge
# endpoints; they must not become writes. They stay in
# acc_vertex, which edge rendering reads, and are dropped here.
writable = [
x.vertex
for x in vertex_list
if not x.lookup_only and x.find is None
]
if writable:
vertex_list_updated = fuse_doc_basis(
writable,
tuple(self.vertex_config.identity_fields(vertex_name)),
)
vertex_list_updated = pick_unique_dict(vertex_list_updated)
assembly_ctx.acc_global[vertex_name] += vertex_list_updated
# A `find` observation carries no primary key: it is written
# onto the vertex its selector finds, so it is merged on those
# fields and kept apart from the documents that are upserted.
attached = self._fuse_attached(vertex_name, vertex_list)
if attached:
assembly_ctx.acc_attached[vertex_name] += attached
for vertex_name, docs in assembly_ctx.acc_attached.items():
assembly_ctx.acc_attached[vertex_name] = pick_unique_dict(docs)
assembly_ctx = add_blank_collections(assembly_ctx, self.vertex_config)
if isinstance(ctx, ActionContext):
ctx.acc_global = assembly_ctx.acc_global
ctx.acc_attached = assembly_ctx.acc_attached
return ctx.acc_global
return assembly_ctx.acc_global
def _fuse_attached(
self, vertex_name: str, vertex_list: list[VertexRep]
) -> list[dict[str, Any]]:
"""The attached documents of one location, merged on their find fields."""
attached = [x for x in vertex_list if x.find is not None and not x.lookup_only]
if not attached:
return []
fields = self._find_fields(vertex_name, attached[0].find)
return fuse_doc_basis([x.vertex for x in attached], fields)
def _find_fields(self, vertex_name: str, selector: str | None) -> tuple[str, ...]:
"""The fields *selector* finds *vertex_name* by, resolved once per pair.
Keyed by the vertex config object, like :meth:`_inverse_pairs`.
"""
vertex_config = self.init_ctx.vertex_config
cached = self.__dict__.get("_find_fields_cache")
if cached is None or cached[0] is not vertex_config:
cached = (vertex_config, {})
self.__dict__["_find_fields_cache"] = cached
by_pair: dict[tuple[str, str | None], tuple[str, ...]] = cached[1]
fields = by_pair.get((vertex_name, selector))
if fields is None:
fields = tuple(vertex_config.match_fields(vertex_name, selector))
by_pair[vertex_name, selector] = fields
return fields
@classmethod
def from_dict(cls, data: dict | list) -> ActorWrapper:
if isinstance(data, list):
return cls(*data)
return cls(**data)
def assemble_tree(
self,
fig_path: Path | str | None = None,
output_format: str = "pdf",
output_dpi: int | None = None,
):
"""Draw this pipeline's actor tree, or return it as a graph.
Delegates to :func:`graflo.plot.plotter.assemble_tree`, which is the
one implementation. ``graflo.plot`` sits above this layer, so the
import is made here rather than at module scope; a missing plotting
extra is reported rather than raised.
Args:
fig_path: Where to write the figure; ``None`` returns the graph.
output_format: Figure format, when writing one.
output_dpi: Raster resolution, for ``png``.
Returns:
``networkx.MultiDiGraph | None``: the tree when *fig_path* is
``None``, otherwise ``None``.
"""
import logging
logger = logging.getLogger(__name__)
try:
from graflo.plot.plotter import assemble_tree
except ImportError as exc:
logger.error("not able to import the plotting stack: %s", exc)
return None
return assemble_tree(
self,
fig_path=fig_path,
output_format=output_format,
output_dpi=output_dpi,
)
def fetch_actors(self, level: int, edges: list) -> tuple[int, type, str, list]:
return self.actor.fetch_actors(level, edges)
def collect_actors(self) -> list[Actor]:
actors = [self.actor]
if isinstance(self.actor, DescendActor):
for descendant in self.actor.descendants:
actors.extend(descendant.collect_actors())
return actors
def find_descendants(
self,
predicate: Callable[[ActorWrapper], bool] | None = None,
*,
actor_type: type[Actor] | None = None,
**attr_in: Any,
) -> list[ActorWrapper]:
if predicate is None:
def _predicate(w: ActorWrapper) -> bool:
if actor_type is not None and not isinstance(w.actor, actor_type):
return False
for attr, allowed in attr_in.items():
if allowed is None:
continue
val = getattr(w.actor, attr, None)
if val not in allowed:
return False
return True
predicate = _predicate
result: list[ActorWrapper] = []
if isinstance(self.actor, DescendActor):
for d in self.actor.descendants:
if predicate(d):
result.append(d)
result.extend(d.find_descendants(predicate=predicate))
return result
def remove_descendants_if(self, predicate: Callable[[ActorWrapper], bool]) -> None:
if isinstance(self.actor, DescendActor):
for d in list(self.actor.descendants):
d.remove_descendants_if(predicate=predicate)
self.actor._descendants[:] = [
d
for d in self.actor.descendants
if not predicate(d)
and not (isinstance(d.actor, DescendActor) and d.count() == 0)
]