"""AST-based extractor for input/output specs from Python source files.
Parses a ``.py`` file using only the standard-library ``ast`` module so
that it works without any third-party dependencies installed (e.g. in CI
environments that lack pydantic, torch, etc.).
The extractor looks for:
1. A function decorated with ``@gyoza_op(...)`` to discover the
``input_model`` and ``output_model`` class names.
2. Class definitions that inherit from ``BaseModel`` whose names match
the discovered model names, extracting their annotated fields.
Base classes are followed so that a model composed from a shared config
object publishes the inherited fields too. Bases declared in the same file
are resolved directly; imported ones are located on disk from the file's
import statements and parsed the same way — never imported, so the
no-third-party-dependency guarantee holds.
"""
from __future__ import annotations
import ast
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from gyoza.models.op_definition import (
InputSpec,
InputSpecs,
OutputSpec,
OutputSpecs,
)
_TYPE_MAP: dict[str, str] = {
"str": "string",
"int": "int",
"float": "float",
"bool": "boolean",
"bytes": "bytes",
"list": "array",
"dict": "object",
"tuple": "array",
"set": "array",
"frozenset": "array",
"List": "array",
"Dict": "object",
"Tuple": "array",
"Set": "array",
"Optional": "optional",
"Any": "any",
}
# Python literal types, as they appear inside ``Literal[...]``.
_LITERAL_TYPES: tuple[tuple[type, str], ...] = (
(bool, "boolean"),
(int, "int"),
(float, "float"),
(str, "string"),
(bytes, "bytes"),
)
# type_str, required, default, description, example
_FieldInfo = tuple[str, bool, Any, str | None, Any]
@dataclass
class _DecoratorInfo:
input_model_name: str | None = None
output_model_name: str | None = None
@dataclass
class _FieldMeta:
"""Metadata extracted from a ``Field(...)`` call."""
default: Any = None
has_default: bool = False
has_default_factory: bool = False
description: str | None = None
example: Any = None
def _is_none(node: ast.expr) -> bool:
"""Return whether an annotation node is the ``None`` literal."""
return isinstance(node, ast.Constant) and node.value is None
def _union_members(node: ast.expr) -> list[ast.expr]:
"""Flatten a ``X | Y | Z`` annotation into its operands."""
if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
return _union_members(node.left) + _union_members(node.right)
return [node]
def _single_type(nodes: list[ast.expr]) -> str:
"""Resolve a group of annotation nodes to one type, or ``"any"``.
``...`` and ``None`` members are ignored, so ``tuple[str, ...]`` and
``str | None`` both narrow to their meaningful type.
Parameters
----------
nodes : list[ast.expr]
Annotation nodes making up a union or a subscript parameter list.
Returns
-------
str
The shared type string, or ``"any"`` when the members disagree.
"""
resolved = {
_resolve_type(node)
for node in nodes
if not _is_none(node)
and not (isinstance(node, ast.Constant) and node.value is Ellipsis)
}
return resolved.pop() if len(resolved) == 1 else "any"
def _literal_type(node: ast.expr) -> str:
"""Resolve ``Literal[...]`` to the type of the values it allows.
Parameters
----------
node : ast.expr
The subscript slice of a ``Literal[...]`` annotation.
Returns
-------
str
Type string shared by the allowed values, or ``"any"`` when mixed.
"""
elts = node.elts if isinstance(node, ast.Tuple) else [node]
names: set[str] = set()
for elt in elts:
if not isinstance(elt, ast.Constant):
return "any"
for python_type, name in _LITERAL_TYPES:
if isinstance(elt.value, python_type):
names.add(name)
break
return names.pop() if len(names) == 1 else "any"
def _resolve_type(node: ast.expr) -> str:
"""Convert an AST annotation node to a simplified type string.
Parameters
----------
node : ast.expr
The annotation AST node.
Returns
-------
str
Simplified type string (e.g. ``"string"``, ``"array[string]"``).
"""
if isinstance(node, ast.Name):
return _TYPE_MAP.get(node.id, node.id)
if isinstance(node, ast.Constant) and isinstance(node.value, str):
return _TYPE_MAP.get(node.value, node.value)
if isinstance(node, ast.Attribute):
return _TYPE_MAP.get(node.attr, node.attr)
if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
return _single_type(_union_members(node))
if isinstance(node, ast.Subscript):
base = _resolve_type(node.value)
slice_node = node.slice
if base == "Annotated":
if isinstance(slice_node, ast.Tuple) and slice_node.elts:
return _resolve_type(slice_node.elts[0])
return _resolve_type(slice_node)
if base == "Literal":
return _literal_type(slice_node)
if base == "object":
return "object"
inner = _single_type(
slice_node.elts if isinstance(slice_node, ast.Tuple) else [slice_node]
)
return inner if base == "optional" else f"{base}[{inner}]"
return "any"
_ARITHMETIC: dict[type[ast.operator], Any] = {
ast.Add: lambda a, b: a + b,
ast.Sub: lambda a, b: a - b,
ast.Mult: lambda a, b: a * b,
ast.Div: lambda a, b: a / b,
ast.Pow: lambda a, b: a**b,
}
def _arithmetic_value(node: ast.expr) -> tuple[bool, Any]:
"""Fold a constant arithmetic expression such as ``1.0 / 10000.0``.
Only numeric operands and the four basic operators are folded, so nothing
from the parsed file is executed.
Parameters
----------
node : ast.expr
Expression to evaluate.
Returns
-------
tuple[bool, Any]
``(True, value)`` when the expression folds to a number.
"""
if not isinstance(node, ast.BinOp):
return False, None
operation = _ARITHMETIC.get(type(node.op))
if operation is None:
return False, None
left_ok, left = _literal_value(node.left)
right_ok, right = _literal_value(node.right)
if not (left_ok and right_ok):
return False, None
if not isinstance(left, (int, float)) or not isinstance(right, (int, float)):
return False, None
if isinstance(left, bool) or isinstance(right, bool):
return False, None
try:
return True, operation(left, right)
except (ArithmeticError, TypeError):
return False, None
def _json_ready(value: Any) -> Any:
"""Convert tuples and sets to lists so a default survives JSON transport."""
if isinstance(value, (tuple, set, frozenset)):
return [_json_ready(item) for item in value]
if isinstance(value, list):
return [_json_ready(item) for item in value]
if isinstance(value, dict):
return {key: _json_ready(item) for key, item in value.items()}
return value
def _literal_value(node: ast.expr) -> tuple[bool, Any]:
"""Return whether *node* is a supported literal and its value.
Covers everything ``ast.literal_eval`` accepts — constants, negative
numbers, lists, tuples, dicts — plus constant arithmetic. Tuples and sets
are returned as lists so the value stays JSON-serialisable.
Parameters
----------
node : ast.expr
AST expression to evaluate as a literal.
Returns
-------
tuple[bool, Any]
``(True, value)`` for supported literals; ``(False, None)`` otherwise.
"""
try:
return True, _json_ready(ast.literal_eval(node))
except (ValueError, TypeError, SyntaxError, MemoryError, RecursionError):
return _arithmetic_value(node)
def _is_field_call(node: ast.expr) -> bool:
"""Return whether *node* is a ``Field(...)`` / ``pydantic.Field(...)`` call."""
if not isinstance(node, ast.Call):
return False
func = node.func
if isinstance(func, ast.Name):
return func.id == "Field"
if isinstance(func, ast.Attribute):
return func.attr == "Field"
return False
def _parse_field_call(node: ast.Call) -> _FieldMeta:
"""Extract default, description and example from a ``Field(...)`` call.
Parameters
----------
node : ast.Call
The ``Field(...)`` call AST node.
Returns
-------
_FieldMeta
Parsed field metadata.
"""
meta = _FieldMeta()
if node.args:
ok, default = _literal_value(node.args[0])
if ok:
meta.default = default
meta.has_default = True
for kw in node.keywords:
if kw.arg == "default":
ok, default = _literal_value(kw.value)
if ok:
meta.default = default
meta.has_default = True
elif kw.arg == "default_factory":
# The value is computed at runtime, so it is optional but has no
# publishable default.
meta.has_default_factory = True
elif kw.arg == "description":
ok, value = _literal_value(kw.value)
if ok and isinstance(value, str):
meta.description = value
elif kw.arg == "example":
ok, value = _literal_value(kw.value)
if ok:
meta.example = value
elif kw.arg == "examples":
ok, value = _literal_value(kw.value)
if ok and isinstance(value, list) and value:
meta.example = value[0]
return meta
def _find_field_in_annotated(annotation: ast.expr) -> _FieldMeta | None:
"""Find a ``Field(...)`` call inside an ``Annotated[...]`` annotation.
Parameters
----------
annotation : ast.expr
Field type annotation AST node.
Returns
-------
_FieldMeta | None
Metadata from the first nested ``Field`` call, if any.
"""
if not isinstance(annotation, ast.Subscript):
return None
base = annotation.value
base_name = None
if isinstance(base, ast.Name):
base_name = base.id
elif isinstance(base, ast.Attribute):
base_name = base.attr
if base_name != "Annotated":
return None
slice_node = annotation.slice
elts = slice_node.elts if isinstance(slice_node, ast.Tuple) else [slice_node]
for elt in elts[1:]:
if _is_field_call(elt) and isinstance(elt, ast.Call):
return _parse_field_call(elt)
return None
def _is_gyoza_op(node: ast.expr) -> bool:
if isinstance(node, ast.Call):
return _is_gyoza_op(node.func)
if isinstance(node, ast.Name):
return node.id in {"gyoza_op", "GyozaOp"}
if isinstance(node, ast.Attribute):
return node.attr in {"gyoza_op", "GyozaOp"}
return False
def _find_decorator_info(tree: ast.Module) -> _DecoratorInfo:
"""Walk the AST to find the first ``@gyoza_op(...)`` decorator.
Parameters
----------
tree : ast.Module
Parsed module AST.
Returns
-------
_DecoratorInfo
Extracted ``input_model`` and ``output_model`` class names.
Raises
------
ValueError
If no ``@gyoza_op`` decorator is found.
"""
for node in ast.walk(tree):
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
continue
for dec in node.decorator_list:
if not (_is_gyoza_op(dec) and isinstance(dec, ast.Call)):
continue
info = _DecoratorInfo()
for kw in dec.keywords:
if kw.arg == "input_model" and isinstance(kw.value, ast.Name):
info.input_model_name = kw.value.id
elif kw.arg == "output_model" and isinstance(kw.value, ast.Name):
info.output_model_name = kw.value.id
return info
msg = "No @gyoza_op decorator found in the file"
raise ValueError(msg)
# Bases that model frameworks provide; following them would pull in
# unrelated class attributes, and they never carry op fields.
_IGNORED_BASES = frozenset(
{
"BaseModel",
"RootModel",
"GenericModel",
"object",
"ABC",
"ABCMeta",
"Generic",
"Protocol",
"TypedDict",
"NamedTuple",
"Enum",
"StrEnum",
"IntEnum",
}
)
# Packages never worth walking into when resolving a base class on disk.
_IGNORED_MODULE_ROOTS = frozenset(
{"pydantic", "typing", "typing_extensions", "abc", "enum", "dataclasses"}
)
@dataclass
class _ModuleContext:
"""A parsed module plus what it needs to resolve its own base classes.
Parameters
----------
tree : ast.Module
Parsed module AST.
path : Path
Location of the module on disk.
roots : tuple[Path, ...]
Directories searched when resolving an imported module to a file.
"""
tree: ast.Module
path: Path
roots: tuple[Path, ...]
from_imports: dict[str, tuple[str, int]] = field(default_factory=dict)
plain_imports: dict[str, str] = field(default_factory=dict)
def __post_init__(self) -> None:
for node in ast.walk(self.tree):
if isinstance(node, ast.ImportFrom):
for alias in node.names:
local = alias.asname or alias.name
self.from_imports[local] = (node.module or "", node.level)
elif isinstance(node, ast.Import):
for alias in node.names:
local = alias.asname or alias.name.split(".")[0]
self.plain_imports[local] = alias.name
def _search_roots(path: Path) -> tuple[Path, ...]:
"""Build the directory list used to resolve imported modules to files.
Covers the module's own package root, the op directory and its ``src``
layout, then ``sys.path`` so installed packages resolve when present.
Parameters
----------
path : Path
Path of the module being parsed.
Returns
-------
tuple[Path, ...]
Search roots, nearest first, deduplicated.
"""
roots: list[Path] = []
package_root = path.parent
while (package_root / "__init__.py").exists():
if package_root.parent == package_root:
break
package_root = package_root.parent
roots.append(package_root)
for candidate in (path.parent, path.parent / "src", package_root / "src"):
roots.append(candidate)
roots.extend(Path(entry) for entry in sys.path if entry)
seen: set[Path] = set()
unique: list[Path] = []
for root in roots:
if root not in seen and root.is_dir():
seen.add(root)
unique.append(root)
return tuple(unique)
def _module_file(
module: str,
level: int,
context: _ModuleContext,
) -> Path | None:
"""Locate the source file of an imported module without importing it.
Parameters
----------
module : str
Dotted module name, empty for a bare relative import.
level : int
Relative-import level; 0 for an absolute import.
context : _ModuleContext
Module the import was written in.
Returns
-------
Path | None
The module's ``.py`` file, or None when it cannot be found.
"""
parts = module.split(".") if module else []
if parts and parts[0] in _IGNORED_MODULE_ROOTS:
return None
if level:
base = context.path.parent
for _ in range(level - 1):
base = base.parent
roots: tuple[Path, ...] = (base,)
else:
roots = context.roots
for root in roots:
target = root.joinpath(*parts) if parts else root
for candidate in (target.with_suffix(".py"), target / "__init__.py"):
if candidate.is_file():
return candidate
return None
def _load_module(
path: Path,
roots: tuple[Path, ...] | None = None,
) -> _ModuleContext | None:
"""Parse a module from disk, returning None when it cannot be read.
Parameters
----------
path : Path
Module file to parse.
roots : tuple[Path, ...] | None
Search roots to reuse; derived from *path* when omitted.
Returns
-------
_ModuleContext | None
Parsed module context, or None on read/syntax failure.
"""
try:
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
except (OSError, SyntaxError, ValueError):
return None
return _ModuleContext(tree=tree, path=path, roots=roots or _search_roots(path))
def _dotted_name(node: ast.expr) -> str | None:
"""Return the dotted source form of a base-class expression."""
if isinstance(node, ast.Name):
return node.id
if isinstance(node, ast.Attribute):
prefix = _dotted_name(node.value)
return f"{prefix}.{node.attr}" if prefix else None
if isinstance(node, ast.Subscript):
return _dotted_name(node.value)
return None
def _find_class(tree: ast.Module, class_name: str) -> ast.ClassDef | None:
"""Return the first class definition named *class_name* in *tree*."""
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef) and node.name == class_name:
return node
return None
def _resolve_base(
node: ast.expr,
context: _ModuleContext,
) -> tuple[ast.ClassDef, _ModuleContext] | None:
"""Resolve a base-class expression to its definition and home module.
Parameters
----------
node : ast.expr
Base-class expression from a ``ClassDef``.
context : _ModuleContext
Module the subclass is declared in.
Returns
-------
tuple[ast.ClassDef, _ModuleContext] | None
The base class and the module holding it, or None when the base is a
framework class or cannot be located on disk.
"""
dotted = _dotted_name(node)
if not dotted:
return None
head, _, tail = dotted.partition(".")
class_name = tail.rsplit(".", 1)[-1] if tail else head
if class_name in _IGNORED_BASES:
return None
if not tail:
local = _find_class(context.tree, class_name)
if local is not None:
return local, context
if head not in context.from_imports:
return None
module, level = context.from_imports[head]
else:
if head not in context.plain_imports:
return None
module, level = context.plain_imports[head], 0
prefix = tail.rsplit(".", 1)[0] if "." in tail else ""
if prefix:
module = f"{module}.{prefix}"
path = _module_file(module, level, context)
if path is None:
return None
# A package may re-export the class; follow the import chain from there.
imported = _load_module(path, context.roots)
if imported is None:
return None
found = _find_class(imported.tree, class_name)
if found is not None:
return found, imported
if class_name in imported.from_imports:
return _resolve_base(ast.Name(id=class_name), imported)
return None
def _extract_class_fields(
tree: ast.Module,
class_name: str,
*,
path: Path | None = None,
) -> dict[str, _FieldInfo]:
"""Extract annotated fields from a class definition and its base classes.
Parameters
----------
tree : ast.Module
Parsed module AST.
class_name : str
Name of the class to extract fields from.
path : Path | None
Location of the module, used to resolve imported base classes.
Returns
-------
dict[str, _FieldInfo]
Mapping of field name to
``(type_str, required, default, description, example)``.
Raises
------
ValueError
If the class is not found.
"""
node = _find_class(tree, class_name)
if node is None:
msg = f"Class '{class_name}' not found in the file"
raise ValueError(msg)
source = path or Path.cwd() / "<memory>"
context = _ModuleContext(tree=tree, path=source, roots=_search_roots(source))
return _collect_fields(node, context, seen=set())
def _collect_fields(
node: ast.ClassDef,
context: _ModuleContext,
seen: set[tuple[str, str]],
) -> dict[str, _FieldInfo]:
"""Collect a class's own fields on top of the ones it inherits.
Bases are walked in reverse declaration order before the class body, so a
subclass redeclaring an inherited field wins — the order Pydantic itself
resolves fields in.
Parameters
----------
node : ast.ClassDef
Class to collect fields from.
context : _ModuleContext
Module holding *node*.
seen : set[tuple[str, str]]
Already-visited ``(module path, class name)`` pairs, guarding against
import cycles and diamond bases.
Returns
-------
dict[str, _FieldInfo]
Field name to extracted metadata, inherited fields first.
"""
marker = (str(context.path), node.name)
if marker in seen:
return {}
seen.add(marker)
fields: dict[str, _FieldInfo] = {}
for base in reversed(node.bases):
resolved = _resolve_base(base, context)
if resolved is not None:
base_node, base_context = resolved
fields.update(_collect_fields(base_node, base_context, seen))
fields.update(_own_fields(node))
return fields
def _own_fields(node: ast.ClassDef) -> dict[str, _FieldInfo]:
"""Extract the annotated fields declared directly in a class body.
Parameters
----------
node : ast.ClassDef
Class whose body is scanned.
Returns
-------
dict[str, _FieldInfo]
Field name to ``(type_str, required, default, description, example)``.
"""
fields: dict[str, _FieldInfo] = {}
for stmt in node.body:
if not isinstance(stmt, ast.AnnAssign) or not isinstance(stmt.target, ast.Name):
continue
name = stmt.target.id
type_str = _resolve_type(stmt.annotation)
description: str | None = None
example: Any = None
default: Any = None
required = True
annotated_meta = _find_field_in_annotated(stmt.annotation)
if annotated_meta is not None:
description = annotated_meta.description
example = annotated_meta.example
if annotated_meta.has_default:
default = annotated_meta.default
required = False
elif annotated_meta.has_default_factory:
required = False
if stmt.value is not None:
if _is_field_call(stmt.value) and isinstance(stmt.value, ast.Call):
meta = _parse_field_call(stmt.value)
if meta.description is not None:
description = meta.description
if meta.example is not None:
example = meta.example
if meta.has_default:
default = meta.default
required = False
elif meta.has_default_factory:
required = False
else:
# Field(...) without default → still required
required = True if default is None else required
elif isinstance(stmt.value, ast.Constant):
default = stmt.value.value
required = False
else:
# Non-literal, non-Field RHS (e.g. function call) → optional, no default
required = False
elif annotated_meta is None:
required = True
fields[name] = (type_str, required, default, description, example)
return fields
