orchestrata_tools.test_generator.testgen.writer

Reference › orchestrata_tools › test_generator › testgen › writer

This module builds indented Python source so the generators avoid string handling. It is imported by branch, structural and generator.

CodeBuilder accumulates lines with an indent depth. line, blank, comment, docstring and todo add content, block and end open and close an indented block, cls and func open class and function definitions (with bases, decorators and async), pytest_raises opens a raises block, and render returns the text. TestFile collects sections for one generated file, with add_section, add_import for standard-library modules the code references, and add_typing_import for names such as cast and Any, then render. The typing registration was added so that whenever a generator emits cast(Any, None) the import lands in the same file. The mock-related scaffolding of the earlier version was dropped, because the project does not mock.

At a glance

Class diagram

Dependency diagram

API

Classes

CodeBuilder

Line 26

Accumulate lines of Python source with indent tracking.

TestFile (dataclass)

Decorators: dataclass · Line 109

Holds all content for one generated test file.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_tools/test_generator/testgen/writer.py

"""writer.py — Indented Python source builder.

Keeps test generators free of string-wrangling.

Restored 2026-08-19 from ``scripts/ast_serializer/testgen`` (deleted
2026-07-09, commit 6a186c9d). Self-contained then and now — the only
change is dropping ``mock_patch``/the conditional ``unittest.mock`` import
block: this project's hard rule is "no mocking, ever" (the only sanctioned
test double is the LLM boundary), so emitting mock-capable scaffolding had
no legitimate use, mocking-config toggle or not. In its place, ``TestFile``
gained ``typing_imports`` — template-defect fix (b): whenever a generator
emits ``cast(Any, None)`` it must add ``from typing import Any, cast`` to
this same file, not leave it to be discovered (and dropped) by a human
copying the pattern later.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from orchestrata_tools.test_generator.testgen.config import ProjectConfig


class CodeBuilder:
    """Accumulate lines of Python source with indent tracking."""

    def __init__(self, indent: int = 4) -> None:
        self._lines: list[str] = []
        self._depth = 0
        self._indent = indent

    # ------------------------------------------------------------------ emit
    def line(self, text: str = "") -> CodeBuilder:
        if text:
            self._lines.append(" " * (self._depth * self._indent) + text)
        else:
            self._lines.append("")
        return self

    def blank(self, n: int = 1) -> CodeBuilder:
        for _ in range(n):
            self._lines.append("")
        return self

    def comment(self, text: str) -> CodeBuilder:
        return self.line(f"# {text}")

    def docstring(self, text: str) -> CodeBuilder:
        text = text.strip()
        if "\n" in text:
            self.line('"""')
            for line in text.splitlines():
                self.line(line)
            self.line('"""')
        else:
            self.line(f'"""{text}"""')
        return self

    def todo(self, hint: str = "") -> CodeBuilder:
        msg = f"TODO: {hint}" if hint else "TODO"
        return self.line(f"# {msg}")

    # ------------------------------------------------------------------ blocks
    def block(self, header: str) -> CodeBuilder:
        self.line(f"{header}:")
        self._depth += 1
        return self

    def end(self) -> CodeBuilder:
        self._depth -= 1
        return self

    # ------------------------------------------------------------------ convenience blocks
    def cls(self, name: str, bases: list[str] | None = None) -> CodeBuilder:
        b = f"({', '.join(bases)})" if bases else ""
        return self.block(f"class {name}{b}")

    def func(
        self,
        name: str,
        args: str = "self",
        decorators: list[str] | None = None,
        is_async: bool = False,
    ) -> CodeBuilder:
        for dec in decorators or []:
            self.line(f"@{dec}")
        kw = "async def" if is_async else "def"
        return self.block(f"{kw} {name}({args})")

    def pytest_raises(self, exc: str) -> CodeBuilder:
        return self.block(f"with pytest.raises({exc})")

    # ------------------------------------------------------------------ output
    def render(self) -> str:
        return "\n".join(self._lines)

    def __str__(self) -> str:
        return self.render()


# ---------------------------------------------------------------------------
# File-level assembler
# ---------------------------------------------------------------------------


@dataclass
class TestFile:
    """Holds all content for one generated test file."""

    module_under_test: str  # dotted module path
    import_lines: list[str] = field(default_factory=list)
    sections: list[str] = field(default_factory=list)  # rendered blocks
    extra_imports: set[str] = field(default_factory=set)
    #: Names to pull from ``typing`` — e.g. {"Any", "cast"}. See
    #: ``add_typing_import`` / template-defect fix (b).
    typing_imports: set[str] = field(default_factory=set)
    config: ProjectConfig | None = None

    def add_section(self, code: CodeBuilder | str) -> None:
        self.sections.append(str(code))

    def add_import(self, module: str) -> None:
        """Register a stdlib module that generated code references."""
        self.extra_imports.add(module)

    def add_typing_import(self, name: str) -> None:
        """Register a ``typing`` name (e.g. "cast", "Any") that generated
        code references — rendered as one ``from typing import ...`` line
        in this same file."""
        self.typing_imports.add(name)

    def render(self) -> str:
        parts: list[str] = []

        # Header
        parts.append('"""')
        parts.append(f"Auto-generated tests for: {self.module_under_test}")
        parts.append("")
        parts.append("Generated by orchestrata_tools.test_generator.testgen")
        parts.append("Edit freely — this file is NOT regenerated automatically.")
        parts.append('"""')
        parts.append("import pytest")

        if self.typing_imports:
            parts.append(f"from typing import {', '.join(sorted(self.typing_imports))}")

        # Extra stdlib imports needed by generated code
        for mod in sorted(self.extra_imports):
            parts.append(f"import {mod}")

        for imp in self.import_lines:
            parts.append(imp)
        parts.append("")
        parts.append("")

        for section in self.sections:
            parts.append(section)
            parts.append("")
            parts.append("")

        return "\n".join(parts).rstrip() + "\n"