Pin Return-Shape Schemas Before You Extract One Function

A messy-repo refactor stays unsafe until return shapes are locked. Lock those shapes with schema tests before any extract. Then change one function and leave everything else frozen.

This article is a procedure, not a personal memoir. Every command below is a proposed local workflow. Adapt file names to your tree before you run anything.

Schema tests are characterization tests for return shapes. They freeze keys, types, and nullability at a boundary. They do not freeze private locals or log text.

Why shape drift beats a green diff

A green unit suite can still hide a contract break. Helpers start returning extra keys or dropped fields. Callers keep working until a serializer or client fails.

String snapshots mix incidental noise with real contract signal. Whitespace, key order, and timestamps pollute golden files.

A JSON Schema oracle cares about types and required keys. That contract is the smallest honest fence for messy modules.

AI-assisted extracts fail in one predictable, repeated way. The model clarifies a key name or wraps a scalar.

The unit tests still pass on mocked internals. The schema test is the first check that sees the leak.

What you pin, and what you refuse to pin

Pin only values that cross a module boundary. Refuse private locals, log lines, and wall-clock fields.

Pin required keys, scalar types, nullability, and closed enums. Leave free-form text out unless a client parses it.

Proposed rule: keep one schema file per public function. Name that file after the function, not after the test. Commit the schema before any model is asked to edit.

Artifact: a schema-lock harness

The harness below is proposed Python, not a recorded run. Replace the messy_mod.build_invoice import with your real seam. Keep fixtures redacted if they resemble customer data.

1. Freeze the working tree

git status --porcelain
git rev-parse HEAD
test -z "$(git status --porcelain)"

Stop if the porcelain output is not empty. A dirty tree makes the later oracle lie. Record the HEAD hash inside the schema commit message.

2. Capture one return value as JSON

# proposed file: tools/capture_shape.py
import inspect
import json
from pathlib import Path

# Replace this import with your boundary function.
from messy_mod import build_invoice


def capture(fn, fixture_path: Path, out_path: Path) -> None:
    payload = json.loads(fixture_path.read_text())
    result = fn(**payload)
    if not isinstance(result, dict):
        raise TypeError("boundary must return a dict for this harness")
    out_path.write_text(json.dumps(result, sort_keys=True, indent=2))
    print(inspect.signature(fn))
    print(out_path)


if __name__ == "__main__":
    capture(
        build_invoice,
        Path("fixtures/build_invoice.input.json"),
        Path("oracles/build_invoice.sample.json"),
    )

Run it once against a known, redacted fixture. Do not pretty-print the sample by hand. Sorted keys keep later review diffs small and boring.

3. Derive a draft schema from the sample

# proposed file: tools/shape_to_schema.py
import json
from pathlib import Path
from typing import Any


def infer(value: Any) -> dict:
    if value is None:
        return {"type": "null"}
    if isinstance(value, bool):
        return {"type": "boolean"}
    if isinstance(value, int) and not isinstance(value, bool):
        return {"type": "integer"}
    if isinstance(value, float):
        return {"type": "number"}
    if isinstance(value, str):
        return {"type": "string"}
    if isinstance(value, list):
        if not value:
            return {"type": "array"}
        return {"type": "array", "items": infer(value[0])}
    if isinstance(value, dict):
        props = {k: infer(v) for k, v in sorted(value.items())}
        required = [k for k, v in value.items() if v is not None]
        return {
            "type": "object",
            "properties": props,
            "required": sorted(required),
            "additionalProperties": False,
        }
    raise TypeError(f"unsupported: {type(value)!r}")


if __name__ == "__main__":
    sample = json.loads(Path("oracles/build_invoice.sample.json").read_text())
    schema = infer(sample)
    Path("oracles/build_invoice.schema.json").write_text(
        json.dumps(schema, indent=2) + "\n"
    )

This inference step is conservative and still incomplete. Empty lists yield arrays with no item schema.

Mixed unions still need a careful human pass. Treat the draft as a checklist, not as truth.

4. Turn the schema into a failing-closed test

# proposed file: tests/test_build_invoice_shape.py
import json
from pathlib import Path

from jsonschema import Draft202012Validator

from messy_mod import build_invoice

SCHEMA = json.loads(Path("oracles/build_invoice.schema.json").read_text())
INPUTS = json.loads(Path("fixtures/build_invoice.input.json").read_text())
VALIDATOR = Draft202012Validator(SCHEMA)


def test_build_invoice_matches_pinned_schema():
    result = build_invoice(**INPUTS)
    errors = sorted(VALIDATOR.iter_errors(result), key=lambda e: e.json_path)
    assert errors == [], [e.message for e in errors]

Run that test before you touch production code. A red test means the capture step was wrong. Do not start the extract until this test is green.

python -m pip install jsonschema pytest
python -m pytest tests/test_build_invoice_shape.py -q

5. Read a schema failure before you extract

Proposed pytest output looks like this when a key leaks:

E   AssertionError: ["Additional properties are not allowed ('tax_hint' was unexpected)"]

Translate that line before you touch the extract. Extra keys mean the helper started publishing internals.

Missing required keys mean a caller will throw later. Type flips from integer to string mean a serializer changed.

Do not widen the schema to silence that failure. Revert the helper instead, then retry a smaller cut.

6. Make the smallest safe change

Extract one private helper from the public function. Do not rename public functions in the same commit.

Do not clean adjacent dead code during this step. Re-run the schema test after the extract lands.

If it fails, restore the module with git checkout. Then re-run the schema test to confirm the revert.

git checkout -- messy_mod.py
python -m pytest tests/test_build_invoice_shape.py -q

Proposed commit split stays strictly two commits long. Stop there and refuse a third bundled change.

1. test: pin build_invoice return schema
2. refactor: extract _line_total from build_invoice

Decision table: which oracle to pin

Signal you care about	Pin this	Do not pin this
Extra or missing keys	required plus additionalProperties: false	Full pretty-printed dumps
Enumerated status strings	enum on that field only	Entire log stream
Numeric type changes	integer versus number	Floating noise beyond the contract
Nested list element shape	items schema	List order if the API is a set
Exception types	A tiny table of error classes	Traceback text

Pick one row from that table per commit. Do not pin all five signals in one commit. Mixed oracles hide which field actually failed later.

Where a free coding model belongs

A model is useful only after the schema file exists. It is a poor first reader of a god module.

Ask it to propose the extract and nothing else. Ask it to keep the public function signature unchanged.

Disclosure: This article was prepared as part of MonkeyCode's product outreach.

MonkeyCode offers free model access and a free server option. Those two facts are the only product claims used here.

No model names, quotas, or hardware details are stated. Those details are not verified for this draft.

Host the pytest runner on that free server after the schema commit. Keep secrets off that machine at all times. Keep fixtures private when they might contain customer data.

A short prompt that stays inside the lock:

Public function: build_invoice
Do not change its signature or return schema.
Schema file: oracles/build_invoice.schema.json
Extract _line_total as a private helper.
Show the diff. Do not edit other functions.

Paste that prompt only after pytest is already green. Reject any diff that touches the committed schema file. Reject any diff that rewrites tests to match new keys.

Failure analysis: four common breaks

1. The sample JSON included a moving timestamp field. The schema froze string on a field that changes. Strip clock fields before you run inference.
2. The flag additionalProperties stayed true on the boundary object. Extra keys then slipped through without a failure. Set that flag false at the boundary object only.
3. The extract changed an exception into a None return. Schema tests stayed green because the happy path matched. Add a second test that covers the error path.
4. The model renamed a required key for supposed clarity. The schema test failed, which is the correct outcome. Revert the rename and do not update the schema.

Limitations

This harness does not prove behavior inside the function. It only freezes the shape that callers already depend on.

It will not catch a wrong total that remains an integer. Pair it with one numeric assertion if money moves.

JSON Schema cannot express every live Python object. Datetimes, decimals, and custom classes need an explicit encoder.

If your boundary returns a dataclass, serialize it once. Pin that encoder output, not the live object graph.

The inference script lies on empty arrays and mixed unions. Review those nodes by hand before you commit.

A bare array type with no items is a hole. Fill the items node before you commit the schema file.

Parallel test runs can race if capture writes a shared sample. Write samples to a unique path per fixture name. Never regenerate schema files inside CI as a side effect.

Who should not use this

Do not use this flow on cryptographic code paths. Shape tests will not catch a weak nonce.

Do not use it when you can still rewrite the public API. If the module has no callers, delete it instead.

Do not point a hosted runner at production payloads. Redact fixtures first or keep the oracle local.

Teams with a stable OpenAPI document already have a contract. Use that document rather than a second schema tree.

Checklist before you merge

1. Confirm the working tree was clean before capture.
2. Commit the schema file in its own commit.
3. Keep additionalProperties false at the boundary object.
4. Strip clock fields and trace fields before inference.
5. Land one extract, one helper, and a green pytest file.
6. Refuse any schema edit inside the refactor commit.

A return-shape schema is a cheap, local fence. Put that fence up before you extract one function. Stop the moment the fence still holds after the cut.