frozen pydantic models

This commit is contained in:
John Lancaster
2026-06-21 12:26:12 -05:00
parent 4320a251f5
commit 9a9432cc55
10 changed files with 309 additions and 119 deletions
@@ -0,0 +1 @@
"""Functions to produce immutable dataclasses representing the document registry."""
@@ -0,0 +1,72 @@
from collections.abc import Generator
from dataclasses import dataclass
from dataclasses import field
from importlib.abc import Traversable
from itertools import starmap
from pathlib import PurePosixPath
from typing import Self
@dataclass(frozen=True, slots=True)
class MarkdownDocument:
"""Represents a loaded markdown document with its content and frontmatter."""
relpath: PurePosixPath
"""The relative path of the document within the package resources."""
content: str = field(repr=False)
"""The raw markdown content of the document."""
frontmatter: str | None = field(repr=False, default=None)
"""The raw YAML frontmatter of the document, if present."""
@classmethod
def from_root(cls, root: Traversable):
"""Recursively load all markdown documents from the root resource."""
mapped = starmap(cls.from_resource, walk_resources(root))
return {d.relpath: d for d in mapped}
@classmethod
def from_resource(cls, relpath: PurePosixPath, resource: Traversable) -> Self:
"""Load a markdown document from a package resource."""
raw = resource.read_text(encoding="utf-8")
frontmatter = get_raw_frontmatter(raw)
return cls(relpath=relpath, content=raw, frontmatter=frontmatter)
@property
def skill_slug(self) -> str | None:
parts = self.relpath.parts
if parts[0] == "skills" and len(parts) >= 3:
return parts[1]
def walk_resources(
node: Traversable,
*,
suffix: str = ".md",
prefix: PurePosixPath | None = None,
) -> Generator[tuple[PurePosixPath, Traversable]]:
"""Recursively yield all resources in node, with their full path."""
prefix = prefix if prefix is not None else PurePosixPath()
for child in sorted(node.iterdir(), key=lambda item: item.name):
relpath = prefix.joinpath(child.name)
if child.is_dir():
yield from walk_resources(child, suffix=suffix, prefix=relpath)
continue
if not child.is_file() or not child.name.lower().endswith(suffix):
continue
yield relpath, child
def get_raw_frontmatter(raw: str) -> str | None:
delimiter = iter(get_frontmatter_delim_idx(raw, delimiter="---"))
try:
start = next(delimiter) + 1
end = next(delimiter)
except StopIteration:
return None
return "\n".join(raw.splitlines()[start:end])
def get_frontmatter_delim_idx(raw: str, *, delimiter: str = "---") -> Generator[int]:
for i, line in enumerate(raw.splitlines()):
if line.strip().startswith(delimiter):
yield i
@@ -0,0 +1,57 @@
from collections.abc import Iterable
from dataclasses import dataclass
from fnmatch import fnmatch
from importlib.resources.abc import Traversable
from itertools import groupby
from itertools import starmap
from typing import Self
from .document import MarkdownDocument
@dataclass(frozen=True, slots=True)
class SkillFilesBundle:
"""Represents a skill and all of its associated markdown files."""
slug: str
skill: MarkdownDocument
references: tuple[MarkdownDocument, ...]
other: tuple[MarkdownDocument, ...]
@classmethod
def from_root(cls, root: Traversable) -> list[Self]:
# Should only be used for testing
return list(cls.from_docs(MarkdownDocument.from_root(root).values()))
@classmethod
def from_docs(cls, docs: Iterable[MarkdownDocument]) -> tuple[Self, ...]:
return tuple(starmap(cls.from_paths, group_skill_paths(docs).items()))
@classmethod
def from_paths(cls, slug: str, paths: set[MarkdownDocument]) -> Self:
skill = next(iter(p for p in paths if p.relpath.name == "SKILL.md"))
groups = {
k: tuple(v)
for k, v in groupby(
paths,
key=lambda p: fnmatch(p.relpath.as_posix(), f"skills/{slug}/references/*.md"),
)
}
references = groups.get(True, ())
other = tuple(p for p in groups.get(False, ()) if p != skill)
return cls(
slug=slug,
skill=skill,
references=references,
other=other,
)
def group_skill_paths(docs: Iterable[MarkdownDocument]) -> dict[str, set[MarkdownDocument]]:
"""Group skills from a list of markdown documents by their skill slug."""
s = sorted(
filter(lambda d: d.skill_slug is not None, docs),
key=lambda d: (d.skill_slug or "", d.relpath.stem),
)
grouped = groupby(s, key=lambda doc: doc.skill_slug)
return {k: set(g) for k, g in grouped if k}