partition schemas and model stubs

.github/workflows/ci.yml

@@ -71,7 +71,7 @@ jobs:
         run: |
           make build-py
 
-      - name: Check native Python codegen output
+      - name: Check Python codegen output
         if: ${{ matrix.python-version == '3.12' }}
         run: |
           make generate-codegen-output-py

codegen-output-py/types/.smithy-python-manifest.json

@@ -1,8 +1,13 @@
 {
   "files": {
-    "src/example_types/_private/schemas.py": "1d708d275049f3f4a86b343c3f38104ef4a7e5bf749a27cd1a490c95c8f9db59",
-    "src/example_types/models.py": "51aadd47e9971f2b44fabd47e81e4b04d7156aa61b64454b39f4b6d47603ab45",
-    "src/example_types/models.pyi": "31c14804662d43c11fc6c61d0a49244e5e606b65aadb774df913f6a26e1b135a"
+    "README.md": "b070865ebb1dc46192ea8a45f0961e1b0f812dfccd069537cfc04311b3327660",
+    "pyproject.toml": "3b4edeff9fe5c4641d8729e49abfa1e7c1936484d0512ad273a3724ebb54f40a",
+    "src/example_types/__init__.py": "27ea4bc5ba93a6fc6e1f11e9784565de24e9249f0d3e77a9e2fb64fa47d11012",
+    "src/example_types/_private/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "src/example_types/_private/schemas.py": "b3c1dd3b1ab8511ef45cca22cf415726aec599edfd457b30cdb22482bcd5b799",
+    "src/example_types/models.py": "262a8c40dce4fb282aca6b17214296db0c4fc1da526a50a3155cde0fe7344b6b",
+    "src/example_types/models.pyi": "36be0dc7d1d9402394e531d782fab2392d55cc60a949a1bcae4c4f469de6cc11",
+    "src/example_types/py.typed": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
   },
   "version": 1
 }

codegen-output-py/types/README.md

@@ -0,0 +1,5 @@
+<!-- Code generated by smithy-python-codegen DO NOT EDIT. -->
+
+## example_types Types Package
+
+Example Python codegen output

codegen-output-py/types/pyproject.toml

@@ -0,0 +1,36 @@
+# Code generated by smithy-python-codegen DO NOT EDIT.
+
+[project]
+name = "example_types"
+version = "1.0.0"
+description = "Example Python codegen output"
+readme = "README.md"
+requires-python = ">=3.12"
+keywords = ["smithy", "example_types"]
+license = {text = "Apache-2.0"}
+dependencies = [
+    "smithy-core~=0.10.0",
+]
+
+[build-system]
+requires = [
+    "hatchling",
+]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = [
+    "src/example_types",
+]
+
+[tool.pyright]
+typeCheckingMode = "strict"
+
+[tool.ruff]
+target-version = "py312"
+
+[tool.ruff.lint]
+ignore = ["F841"]
+
+[tool.ruff.format]
+skip-magic-trailing-comma = true

codegen-output-py/types/src/example_types/__init__.py

@@ -0,0 +1 @@
+__version__: str = "1.0.0"

codegen-output-py/types/src/example_types/_private/schemas.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from smithy_core._schema_compact import StaticSchema
 
 n0 = "example"

codegen-output-py/types/src/example_types/models.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from smithy_core._dynamic_models import ModelRegistry
 from smithy_core.exceptions import ModeledError
 
@@ -43,7 +41,7 @@ class ServiceError(ModeledError):
 }
 
 _MODEL_NAMES = (*tuple(_MODEL_INDEX), "Unit")
-__all__ = ("ServiceError", *_MODEL_NAMES)
+__all__ = ("ServiceError", *_MODEL_NAMES)  # pyright: ignore[reportUnsupportedDunderAll]
 
 
 _registry = ModelRegistry(

codegen-output-py/types/src/example_types/models.pyi

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from dataclasses import dataclass
 from enum import IntEnum, StrEnum
 from typing import Self

designs/codegen/index.md

@@ -65,3 +65,4 @@ behavior of generated packages.
 * [Python Symbols](symbols.md)
 * [Python Writer](writer.md)
 * [Standalone Types Generation](types.md)
+* [Code Generation Partitioning](partitioning.md)

designs/codegen/partitioning.md

@@ -0,0 +1,163 @@
+# Code Generation Partitioning
+
+## Abstract
+
+Large standalone types packages split compact schemas and model stubs across
+private modules. Stable facade modules preserve generated imports while
+reducing generator memory and static analysis time.
+
+## Motivation
+
+An EC2 types package contains 4,065 selected shapes. Emitting its schemas and
+model declarations as two large files makes each file the unit of work for
+rendering, staging, parsing, and static analysis.
+
+`FileManifest` releases each rendered file after staging it, so the largest file
+sets the lower bound for source memory during publication. Pyright analyzes the
+model stub as one large declaration file. Lazy schema hydration and model
+construction reduce runtime work but do not reduce these file-level costs.
+
+Private partitions bound the work associated with one generated file. Stable
+facades preserve generated imports and lazy runtime behavior while tools process
+the smaller partition files.
+
+## Settings
+
+`TypesSettings` owns both partition targets:
+
+```python
+@dataclass(frozen=True, slots=True)
+class TypesSettings(ArtifactSettings):
+    selector: Selector = field(default_factory=Selector.identity)
+    generate_inputs_and_outputs: bool = False
+    schema_partition_size: int | None = 250
+    stub_partition_size: int | None = 250
+```
+
+The CLI uses these defaults. Python callers may set either target to a positive
+integer or `None`. A value of `None` emits that surface as one file.
+
+Partitioning starts when the number of planned items exceeds its target. A
+package with 250 schema shapes and 250 model declarations keeps the single-file
+layout.
+
+## Output Layout
+
+Partitioned packages retain the stable module paths used by unpartitioned
+packages.
+
+| Surface | Facade | Partitions |
+| --- | --- | --- |
+| Compact schemas | `<package>._private.schemas` | `<package>._private.schemas_N` |
+| Model stubs | `<package>.models.pyi` | `<package>._private.models_N.pyi` |
+| Runtime models | `<package>.models` | Single module |
+
+### Schema File Layout
+
+```text
+src/<package>/
+├── models.py
+└── _private/
+    ├── __init__.py
+    ├── schemas.py
+    ├── schemas_0.py
+    ├── schemas_1.py
+    └── schemas_N.py
+```
+
+`models.py` imports the `schemas.py` facade. The facade maps schema constants to
+partition indexes and loads the matching `schemas_N.py` module on first access.
+Each partition contains compact schema tuples.
+
+### Stub File Layout
+
+```text
+src/<package>/
+├── models.py
+├── models.pyi
+└── _private/
+    ├── __init__.py
+    ├── models_0.pyi
+    ├── models_1.pyi
+    └── models_N.pyi
+```
+
+`models.pyi` is the public typing facade. It imports declarations from each
+private stub partition. `models.py` remains the runtime module.
+
+`FileManifest` registers the facade and partition files before publication.
+Its ownership rules remove stale, unmodified partition files after a model
+shrinks or a target changes.
+
+Files under `<package>._private` remain private implementation details.
+
+## Schema Partitions
+
+`SchemaPlan` provides schema shapes in deterministic emission order and groups
+each recursive strongly connected component. A recursive group stays in one
+partition so its schemas can use direct local references.
+
+The schema partitioner:
+
+1. Packs recursive groups in plan order using the configured shape target.
+2. Uses the resulting partition count for the final layout.
+3. Estimates the emitted source size of each group.
+4. Selects contiguous boundaries that balance estimated source size.
+
+The estimate includes shape and member identifiers, schema references, and
+trait values. It uses plan metadata instead of rendering schemas a second time.
+The target determines the partition count. Final partitions may exceed the
+target to balance source size, and recursive groups remain indivisible.
+
+The facade maps each schema constant to a partition index. Its `__getattr__`
+imports the selected partition, reads the constant, and caches the result in
+the facade module. `__dir__` includes every planned schema constant.
+
+Partition modules route cross-partition references through the facade. This
+keeps dependencies lazy and avoids direct import cycles between partitions.
+Hydration remains lazy through `smithy_core._schema_compact`.
+
+## Stub Partitions
+
+`ModelPlan` provides model declarations in deterministic order. The stub
+partitioner writes contiguous groups up to the configured target.
+
+The public `models.pyi` facade imports every declaration from its private stub
+partition and defines the same `__all__` as `models.py`. Private stubs reference
+generated declarations through the public `models` module. This keeps their
+type references independent of private partition boundaries.
+
+Runtime `models.py` remains one lazy registry module. Runtime model construction
+and class identity stay unchanged.
+
+## Compatibility
+
+Partitioning preserves:
+
+* Public model module paths and runtime schema lookups.
+* Model names, constructor signatures, and type annotations.
+* Public model `__all__`, explicit imports, star imports, and `dir()` results.
+* Lazy schema hydration and model construction.
+* Schema identity, recursive references, and serde behavior.
+
+Partition indexes follow generation-plan order. Model changes can move
+declarations between private files while the public facades remain stable.
+
+## Validation
+
+Tests cover packages below and above each target, recursive schema groups,
+cross-partition references, facade caching, strict Pyright, Ruff, and
+deterministic output.
+
+The final x86 EC2 benchmark compared unpartitioned output with 250-item schema
+and stub partitions:
+
+| Metric | Unpartitioned | Partitioned | Change |
+| --- | ---: | ---: | ---: |
+| Generation time | 3.557 s | 3.555 s | -0.1% |
+| Peak RSS | 264.5 MiB | 122.4 MiB | -53.7% |
+| Peak Python allocation | 222.6 MiB | 84.7 MiB | -61.9% |
+| Generated bytes | 4.78 MiB | 5.33 MiB | +11.6% |
+| Model import | 50.6 ms | 40.4 ms | -20.2% |
+| First model access | 12.3 ms | 22.3 ms | +10.0 ms |
+| Strict Pyright | 81.98 s | 15.25 s | -81.4% |

designs/codegen/types.md

@@ -100,19 +100,23 @@ class ArtifactSettings:
 class TypesSettings(ArtifactSettings):
     selector: Selector = field(default_factory=Selector.identity)
     generate_inputs_and_outputs: bool = False
+    schema_partition_size: int | None = 250
+    stub_partition_size: int | None = 250
 ```
 
 The resolved service remains on `Selection`. It is not duplicated in settings.
 The CLI subcommand identifies the artifact, so settings do not duplicate that
 value. Artifact-specific settings extend the shared frozen dataclass.
+Partition targets and their output layout are defined in
+[Code Generation Partitioning](partitioning.md).
 
 The `types` command accepts:
 
 * `--package`
 * `--package-version`
 * `--package-description`
 * `--selector`
-* `--generate-inputs-and-outputs`
+* `--include-operation-io`
 
 `package` and `package-version` are required. The package identifies the
 generated import package and follows the validation rules in `SymbolProvider`.
@@ -129,7 +133,7 @@ Types generation never infers a service. The options apply in this order:
 2. `--selector` narrows those declaration shapes. The default identity selector
    keeps every candidate.
 3. Input and output structures are removed as roots unless
-   `--generate-inputs-and-outputs` is set.
+   `--include-operation-io` is set.
 4. Member-target dependencies of the remaining declarations are added
    transitively.
 
@@ -382,16 +386,17 @@ model names:
 
 ```python
 _MODEL_NAMES = tuple(_MODEL_INDEX)
-__all__ = ("ServiceError", *_MODEL_NAMES)
+__all__ = ("ServiceError", *_MODEL_NAMES)  # pyright: ignore[reportUnsupportedDunderAll]
 
 
 def __dir__() -> list[str]:
     return sorted({*globals(), *__all__})
 ```
 
-This keeps Ruff from reporting `F822` for lazy exports. Star imports read
-`__all__` and resolve each model through module `__getattr__`. Computing
-`__dir__` from the index does not construct model classes.
+This keeps Ruff from reporting `F822` for lazy exports. The targeted Pyright
+suppression replaces excluding the complete runtime module from analysis. Star
+imports read `__all__` and resolve each model through module `__getattr__`.
+Computing `__dir__` from the index does not construct model classes.
 
 Model member access uses the schema's dense member indexes where positional
 dispatch is required. Name maps remain dictionaries where lookup is by a wire
@@ -415,26 +420,12 @@ Package setup renders the dependency set collected during generation.
 
 ## Large Models
 
-Schema and stub partitioning must not change public imports. The output API and
-generation plan allow private content to move behind stable facade modules.
+Large standalone types packages partition compact schemas and model stubs
+behind stable facade modules. Schema partitions use source-size weighting and
+keep recursive components together. Stub partitions use declaration counts.
 
-smithy-java partitions schema classes after a model reaches 500 schema fields
-and then targets about 50 KB of generated schema content per class. smithy-php
-uses a centralized lazy schema registry. Python can combine those approaches:
-
-* Keep `<package>._private.schemas` as the stable facade.
-* Place compact schema tuples in private partition modules.
-* Keep every recursive strongly connected component in one partition.
-* Route schema constant access through generated indexes.
-* Partition large type stubs behind the public `models.pyi` facade.
-
-Python partitioning thresholds must be based on Python measurements. Relevant inputs include writer render time, peak generator memory, generated import time, and type
-checker time.
-
-Partitioning bounds the syntax tree analyzed by `PythonWriter` and prevents a
-single EC2-sized source file from dominating codegen memory. Measurements
-establish the partitioning thresholds. Facade modules preserve generated
-imports.
+[Code Generation Partitioning](partitioning.md) defines the output layout,
+thresholds, and compatibility requirements.
 
 ## Performance