"""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 _module_constant(
name: str,
context: _ModuleContext,
seen: set[tuple[str, str]],
) -> tuple[bool, Any]:
"""Read a module-level constant, following the import that brought it in.
Ops name their defaults instead of repeating them — ``MVT_EXTENT`` rather
than ``65536`` — so the value lives one module away from the field that
uses it. Reading it here is what lets the published spec carry the same
number the op would have applied itself.
Parameters
----------
name : str
Constant name as written at the use site.
context : _ModuleContext
Module to look in first.
seen : set[tuple[str, str]]
Already-visited ``(module path, constant name)`` pairs, guarding
against re-export cycles.
Returns
-------
tuple[bool, Any]
``(True, value)`` when the constant resolves to a literal.
"""
marker = (str(context.path), name)
if marker in seen:
return False, None
seen.add(marker)
for stmt in context.tree.body:
if isinstance(stmt, ast.Assign):
if any(
isinstance(target, ast.Name) and target.id == name
for target in stmt.targets
):
return _resolved_value(stmt.value, context, seen)
elif isinstance(stmt, ast.AnnAssign):
target = stmt.target
if (
isinstance(target, ast.Name)
and target.id == name
and stmt.value is not None
):
return _resolved_value(stmt.value, context, seen)
origin = context.from_imports.get(name)
if origin is None:
return False, None
module, level = origin
path = _module_file(module, level, context)
if path is None:
return False, None
imported = _load_module(path, context.roots)
if imported is None:
return False, None
return _module_constant(name, imported, seen)
def _module_attribute(
node: ast.Attribute,
context: _ModuleContext,
seen: set[tuple[str, str]],
) -> tuple[bool, Any]:
"""Read ``module.CONSTANT`` by locating *module* on disk.
Parameters
----------
node : ast.Attribute
Attribute access whose value is the imported module.
context : _ModuleContext
Module the access was written in.
seen : set[tuple[str, str]]
Already-visited ``(module path, constant name)`` pairs.
Returns
-------
tuple[bool, Any]
``(True, value)`` when the attribute resolves to a literal.
"""
holder = node.value
if not isinstance(holder, ast.Name):
return False, None
module = context.plain_imports.get(holder.id)
level = 0
if module is None:
origin = context.from_imports.get(holder.id)
if origin is None:
return False, None
parent, level = origin
module = f"{parent}.{holder.id}" if parent else holder.id
path = _module_file(module, level, context)
if path is None:
return False, None
imported = _load_module(path, context.roots)
if imported is None:
return False, None
return _module_constant(node.attr, imported, seen)
def _resolved_value(
node: ast.expr,
context: _ModuleContext | None = None,
seen: set[tuple[str, str]] | None = None,
) -> tuple[bool, Any]:
"""Evaluate *node* as a literal, following named constants when possible.
Parameters
----------
node : ast.expr
AST expression to evaluate.
context : _ModuleContext | None
Module the expression was written in; without one only literals
resolve.
seen : set[tuple[str, str]] | None
Cycle guard shared across a single resolution.
Returns
-------
tuple[bool, Any]
``(True, value)`` when the expression resolves to a literal.
"""
ok, value = _literal_value(node)
if ok or context is None:
return ok, value
tracked = set() if seen is None else seen
if isinstance(node, ast.Name):
return _module_constant(node.id, context, tracked)
if isinstance(node, ast.Attribute):
return _module_attribute(node, context, tracked)
return False, None
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,
context: _ModuleContext | None = None,
) -> _FieldMeta:
"""Extract default, description and example from a ``Field(...)`` call.
A ``default=`` makes the field optional whether or not its value can be
read here. Requiring what we merely failed to parse would reject every
caller for an input the op fills in on its own.
Parameters
----------
node : ast.Call
The ``Field(...)`` call AST node.
context : _ModuleContext | None
Module the call was written in, used to follow named constants.
Returns
-------
_FieldMeta
Parsed field metadata.
"""
meta = _FieldMeta()
if node.args:
meta.has_default = True
ok, default = _resolved_value(node.args[0], context)
if ok:
meta.default = default
for kw in node.keywords:
if kw.arg == "default":
meta.has_default = True
ok, default = _resolved_value(kw.value, context)
if ok:
meta.default = default
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 = _resolved_value(kw.value, context)
if ok and isinstance(value, str):
meta.description = value
elif kw.arg == "example":
ok, value = _resolved_value(kw.value, context)
if ok:
meta.example = value
elif kw.arg == "examples":
ok, value = _resolved_value(kw.value, context)
if ok and isinstance(value, list) and value:
meta.example = value[0]
return meta
def _find_field_in_annotated(
annotation: ast.expr,
context: _ModuleContext | None = None,
) -> _FieldMeta | None:
"""Find a ``Field(...)`` call inside an ``Annotated[...]`` annotation.
Parameters
----------
annotation : ast.expr
Field type annotation AST node.
context : _ModuleContext | None
Module the annotation was written in, used to follow named constants.
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, context)
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, context))
return fields
def _own_fields(
node: ast.ClassDef,
context: _ModuleContext | None = None,
) -> dict[str, _FieldInfo]:
"""Extract the annotated fields declared directly in a class body.
Parameters
----------
node : ast.ClassDef
Class whose body is scanned.
context : _ModuleContext | None
Module holding *node*, used to follow named constants.
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, context)
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, context)
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
else:
ok, value = _resolved_value(stmt.value, context)
if ok:
default = value
# A non-literal RHS (e.g. a function call) is still a default,
# just not one we can publish.
required = False
elif annotated_meta is None:
required = True
fields[name] = (type_str, required, default, description, example)
return fields
