from dataclasses import dataclass, field
from typing import TYPE_CHECKING
from pydantic import BaseModel
if TYPE_CHECKING:
from collections.abc import Iterator
from pydantic.fields import FieldInfo
from codegen.store import NodePath
[docs]
@dataclass(frozen=True)
class Scoped:
name: str
[docs]
@dataclass(frozen=True)
class Scope:
name: str
config_key: str
parent: Scope | None = None
resolve_path: tuple[str, ...] = field(default=())
def __post_init__(self) -> None:
if self.parent is None and self.name != "project":
msg = (
f"Scope {self.name!r} has no parent; only the root "
f"'project' scope may be parentless"
)
raise ValueError(msg)
PROJECT = Scope(name="project", config_key="")
[docs]
class ScopeTree(tuple[Scope, ...]):
__slots__ = ()
def children_of(self, parent: Scope) -> list[Scope]:
return [scope for scope in self if scope.parent is parent]
def scope_for(self, instance_id: NodePath) -> Scope:
parts = instance_id.parts
if parts and parts[0] != "project":
msg = f"Instance id {instance_id!r} must start with 'project'"
raise ValueError(msg)
current = PROJECT
for i in range(1, len(parts), 2):
config_key = parts[i]
try:
current = next(
scope
for scope in self
if scope.parent is current
and scope.config_key == config_key
)
except StopIteration as exc:
msg = (
f"Instance id {instance_id!r} references config_key "
f"{config_key!r}, which is not a child of "
f"{current.name!r}"
)
raise ValueError(msg) from exc
return current
[docs]
def discover_scopes(
config_cls: type[BaseModel],
) -> ScopeTree:
return ScopeTree((PROJECT, *_discover(config_cls, PROJECT, (), set())))
def _discover(
cls: type[BaseModel],
parent: Scope,
prefix: tuple[str, ...],
seen: set[type[BaseModel]],
) -> Iterator[Scope]:
for name, info in cls.model_fields.items():
marker = next(
(
potential_marker
for potential_marker in info.metadata
if isinstance(potential_marker, Scoped)
),
None,
)
if marker:
item_cls = _extract_base_model_from_scoped(cls, name, info)
child = Scope(
name=marker.name,
config_key=name,
parent=parent,
resolve_path=(*prefix, name),
)
yield child
if item_cls not in seen:
seen |= {item_cls}
yield from _discover(item_cls, child, (), seen)
elif isinstance(info.annotation, type) and issubclass(
info.annotation, BaseModel
):
yield from _discover(info.annotation, parent, (*prefix, name), seen)
def _extract_base_model_from_scoped(
cls: type[BaseModel],
name: str,
info: FieldInfo,
) -> type[BaseModel]:
annotation = info.annotation
if getattr(annotation, "__origin__", None) is list:
args: tuple[type, ...] = getattr(annotation, "__args__", ())
item = args[0] if args else None
if isinstance(item, type) and issubclass(item, BaseModel):
return item
msg = (
f"{cls.__name__}.{name} is annotated Scoped() but its "
f"type is not list[BaseModel]: {annotation!r}"
)
raise TypeError(msg)