levain-cli 0.1.0__py3-none-any.whl

levain_recipe/codegen.py

@@ -0,0 +1,114 @@
+"""Code generation for recipe packages.
+
+Generates a Python ``GraphState`` TypedDict from the recipe's
+spec.  Node functions type their state parameter as this class,
+so mypy catches any key not declared in the spec.
+"""
+
+from __future__ import annotations
+
+from .models import SPEC_TYPES, StateField, ValidatedSpec
+from .reducers import REDUCERS
+
+_HEADER_BODY = '''"""Generated GraphState TypedDict.
+
+Do not edit by hand — regenerate via::
+
+    levain generate-state <dir>
+"""
+'''
+
+
+def generate_state_module(spec: ValidatedSpec) -> str:
+    """Return Python source for the recipe's state TypedDict."""
+    reducer_names = sorted(
+        {v.reducer for v in spec.variables if v.reducer is not None}
+    )
+    needs_annotated = bool(reducer_names) or any(
+        v.description for v in spec.variables
+    )
+
+    lines: list[str] = [_HEADER_BODY]
+
+    if reducer_names:
+        # Merge ``Any`` into the typing import when any reducer's
+        # element type references it (e.g. ``dict[str, Any]``);
+        # without this the generated file would use ``Any``
+        # without importing it.
+        typing_imports = ["Annotated", "TypedDict"]
+        if any("Any" in REDUCERS[n].element_type_name for n in reducer_names):
+            typing_imports.insert(1, "Any")
+        lines.append(f"from typing import {', '.join(typing_imports)}")
+        lines.append("")
+        # Group reducer imports by module so ruff/isort won't
+        # split them and cause CI drift.
+        by_module: dict[str, list[str]] = {}
+        for name in reducer_names:
+            raw = REDUCERS[name].import_line
+            mod, _, names = raw.partition(" import ")
+            by_module.setdefault(mod, []).append(names)
+        for mod, groups in by_module.items():
+            seen: set[str] = set()
+            merged_names: list[str] = []
+            for group in groups:
+                for n in group.split(","):
+                    name = n.strip()
+                    if name not in seen:
+                        seen.add(name)
+                        merged_names.append(name)
+            lines.append(f"{mod} import {', '.join(merged_names)}")
+    elif needs_annotated:
+        lines.append("from typing import Annotated, TypedDict")
+    else:
+        lines.append("from typing import TypedDict")
+    lines.append("")
+    lines.append("")
+    lines.append("class GraphState(TypedDict, total=False):")
+    lines.append(f'    """State fields for the {spec.module_name} recipe.')
+    lines.append("")
+    lines.append("    ``total=False`` allows node functions to return partial")
+    lines.append("    updates — only the keys they actually change.")
+    lines.append('    """')
+    lines.append("")
+
+    if spec.variables:
+        for var in spec.variables:
+            lines.append(f"    {var.name}: {_py_type_for(var)}")
+    else:
+        # TypedDict body must not be empty.
+        lines.append("    pass")
+
+    return "\n".join(lines) + "\n"
+
+
+def _py_type_for(var: StateField) -> str:
+    """Map a spec variable type to a Python type name.
+
+    For list-typed fields with a blessed reducer, emit
+    ``Annotated[list[<element>], <reducer>]`` so LangGraph picks
+    the reducer up from the TypedDict annotation. When the variable
+    carries a description it's appended as the final metadata arg
+    so ``get_type_hints(..., include_extras=True)`` recovers it.
+    """
+    py_type = SPEC_TYPES.get(var.type)
+    if py_type is None:
+        raise ValueError(
+            f"Unknown type {var.type!r} for variable {var.name!r}"
+        )
+    desc_literal = repr(var.description) if var.description else None
+
+    if var.reducer is not None:
+        spec = REDUCERS.get(var.reducer)
+        if spec is None:
+            raise ValueError(
+                f"Unknown reducer {var.reducer!r} for variable {var.name!r}"
+            )
+        metadata = [spec.name]
+        if desc_literal is not None:
+            metadata.append(desc_literal)
+        return (
+            f"Annotated[list[{spec.element_type_name}], {', '.join(metadata)}]"
+        )
+    if desc_literal is not None:
+        return f"Annotated[{py_type.__name__}, {desc_literal}]"
+    return py_type.__name__