Source code for codegen.operation

import importlib.metadata
from dataclasses import dataclass, field
from graphlib import CycleError, TopologicalSorter
from typing import Any, NamedTuple

from pydantic import BaseModel

_META_ATTR = "_operation_meta"




@dataclass(frozen=True)
class OperationMeta:
    name: str
    scope: str
    requires: tuple[str, ...] = ()
    after_children: bool = False
    dispatch_on: str | None = None
    match_value: str | None = None





class EmptyOptions(BaseModel):
    pass



class OperationEntry(NamedTuple):
    meta: OperationMeta
    cls: type




@dataclass
class OperationRegistry:
    entries: list[OperationEntry] = field(default_factory=list)

    def register(self, meta: OperationMeta, cls: type) -> None:
        self.entries.append(OperationEntry(meta=meta, cls=cls))

    def validate_scopes(self, known: set[str]) -> None:
        for entry in self.entries:
            if entry.meta.scope not in known:
                msg = (
                    f"Operation '{entry.meta.name}' targets "
                    f"scope '{entry.meta.scope}' which was not "
                    f"discovered from the config"
                )

                raise ValueError(msg)

    def sorted_by_scope(self) -> dict[str, list[OperationEntry]]:
        buckets: dict[str, list[OperationEntry]] = {}

        for entry in self.entries:
            buckets.setdefault(entry.meta.scope, []).append(entry)

        return {name: _topo_sort(ops) for name, ops in buckets.items()}





def operation(
    name: str,
    *,
    scope: str,
    requires: list[str] | None = None,
    after_children: bool = False,
    dispatch_on: str | None = None,
    registry: OperationRegistry | None = None,
    match_value: str | None = None,
) -> Any:  # noqa: ANN401
    reqs: tuple[str, ...] = tuple(requires or [])
    meta = OperationMeta(
        name=name,
        scope=scope,
        requires=reqs,
        after_children=after_children,
        dispatch_on=dispatch_on,
        match_value=match_value,
    )

    def decorator(cls: type) -> type:
        setattr(cls, _META_ATTR, meta)

        if not hasattr(cls, "Options"):
            cls.Options = EmptyOptions

        if registry is not None:
            registry.register(meta, cls)

        return cls

    return decorator





def load_registry(entry_point_group: str) -> OperationRegistry:
    registry = OperationRegistry()

    for entry_point in importlib.metadata.entry_points(group=entry_point_group):
        cls = entry_point.load()
        meta = getattr(cls, _META_ATTR, None)

        if meta is None:
            msg = (
                f"Entry point {entry_point.name!r} in group "
                f"{entry_point_group!r} loaded {cls!r}, which is "
                f"not an @operation-decorated class"
            )
            raise TypeError(msg)

        registry.register(meta, cls)

    return registry



def _topo_sort(entries: list[OperationEntry]) -> list[OperationEntry]:
    ordered = sorted(entries, key=lambda entry: entry.meta.name)

    graph: dict[str, set[str]] = {
        entry.meta.name: set(entry.meta.requires) for entry in ordered
    }

    _validate_requires(graph)

    by_name: dict[str, OperationEntry] = {
        entry.meta.name: entry for entry in ordered
    }

    try:
        return [
            by_name[name] for name in TopologicalSorter(graph).static_order()
        ]

    except CycleError as exc:
        msg = "Cycle detected in operation dependencies"
        raise ValueError(msg) from exc


def _validate_requires(graph: dict[str, set[str]]) -> None:
    for name, deps in graph.items():
        for dependency in deps:
            if dependency not in graph:
                msg = (
                    f"Operation '{name}' requires '{dependency}', "
                    f"which is not registered"
                )

                raise ValueError(msg)