From ca23f2f1c66c6a4668c3a2f6e649cd154cad2b26 Mon Sep 17 00:00:00 2001 From: Sean O'Brien <60306702+stobrien89@users.noreply.github.com> Date: Sun, 25 Oct 2026 18:57:01 -0400 Subject: [PATCH] partition schemas and model stubs --- designs/codegen/index.md | 1 + designs/codegen/partitioning.md | 163 +++++++++++ designs/codegen/types.md | 28 +- .../src/smithy_python/generation.py | 9 + .../smithy-python/src/smithy_python/models.py | 124 ++++++++- .../src/smithy_python/schemas.py | 260 ++++++++++++++++-- .../tests/unit/test_generation.py | 9 + .../smithy-python/tests/unit/test_models.py | 85 +++++- .../smithy-python/tests/unit/test_schemas.py | 167 ++++++++++- 9 files changed, 792 insertions(+), 54 deletions(-) create mode 100644 designs/codegen/partitioning.md diff --git a/designs/codegen/index.md b/designs/codegen/index.md index 5455a37b9..81d968097 100644 --- a/designs/codegen/index.md +++ b/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) diff --git a/designs/codegen/partitioning.md b/designs/codegen/partitioning.md new file mode 100644 index 000000000..e4f0c8266 --- /dev/null +++ b/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 | `._private.schemas` | `._private.schemas_N` | +| Model stubs | `.models.pyi` | `._private.models_N.pyi` | +| Runtime models | `.models` | Single module | + +### Schema File Layout + +```text +src// +├── 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// +├── 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 `._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% | diff --git a/designs/codegen/types.md b/designs/codegen/types.md index 58e07fa24..964c4b6f8 100644 --- a/designs/codegen/types.md +++ b/designs/codegen/types.md @@ -100,11 +100,15 @@ 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: @@ -416,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 `._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 diff --git a/packages/smithy-python/src/smithy_python/generation.py b/packages/smithy-python/src/smithy_python/generation.py index ce75e1e19..b82369040 100644 --- a/packages/smithy-python/src/smithy_python/generation.py +++ b/packages/smithy-python/src/smithy_python/generation.py @@ -44,6 +44,15 @@ class TypesSettings(ArtifactSettings): selector: _selection.Selector = field(default_factory=_selection.Selector.identity) generate_inputs_and_outputs: bool = False + schema_partition_size: int | None = 250 + stub_partition_size: int | None = 250 + + def __post_init__(self) -> None: + if any( + size is not None and size < 1 + for size in (self.schema_partition_size, self.stub_partition_size) + ): + raise ValueError("Partition sizes must be positive") @dataclass(frozen=True, slots=True) diff --git a/packages/smithy-python/src/smithy_python/models.py b/packages/smithy-python/src/smithy_python/models.py index 29be7e6ea..a9c3e8b57 100644 --- a/packages/smithy-python/src/smithy_python/models.py +++ b/packages/smithy-python/src/smithy_python/models.py @@ -36,12 +36,31 @@ def generate(self, context: GenerationContext[TypesSettings]) -> None: ), context.plan.models, ) - _StubEmitter( + facade = _StubEmitter( context.files.python_module( f"src/{directory}/models.pyi", f"{package}.models" ), context, - ).write() + ) + shapes = context.plan.models.shapes + size = context.settings.stub_partition_size + if size is None or len(shapes) <= size: + facade.write() + return + + partitions = tuple( + shapes[index : index + size] for index in range(0, len(shapes), size) + ) + facade.write_facade(partitions) + for index, shapes in enumerate(partitions): + _StubEmitter( + context.files.python_module( + f"src/{directory}/_private/models_{index}.pyi", + f"{package}._private.models_{index}", + ), + context, + external_bases=True, + ).write_partition(shapes) def _core_reference(name: str, module: str) -> _symbols.TypeReference: @@ -165,12 +184,25 @@ def __init__( self, writer: PythonWriter, context: GenerationContext[TypesSettings], + *, + external_bases: bool = False, ) -> None: self.writer = writer self.model = context.model self.plan = context.plan.models self.symbols = context.symbols self.module = f"{context.settings.package}.models" + self.external_bases = external_bases + self.serde_shape: str | _symbols.TypeReference = ( + _symbols.TypeReference("_SerdeShape", "..models") + if external_bases + else "_SerdeShape" + ) + self.service_error: str | _symbols.TypeReference = ( + _symbols.TypeReference("ServiceError", "..models") + if external_bases + else "ServiceError" + ) operations = tuple(context.model.iter_shapes(_model.OperationShape)) self.inputs = {operation.input for operation in operations} | { shape.id for shape in self.plan.shapes if shape.has_trait("input") @@ -180,6 +212,30 @@ def __init__( } def write(self) -> None: + """Write the complete unpartitioned public stub.""" + self._write_preamble() + self._write_shapes(self.plan.shapes) + self._write_unit() + + def write_facade(self, partitions: tuple[tuple[_model.Shape, ...], ...]) -> None: + """Write the public stub facade and re-export private declarations.""" + writer = self.writer + for index, shapes in enumerate(partitions): + writer.line(f"from ._private.models_{index} import (") + with writer.indent(): + for shape in shapes: + for name in _stub_shape_names(self.plan, shape): + writer.line(name, ",") + writer.line(")") + writer.line() + self._write_preamble() + self._write_unit() + + def write_partition(self, shapes: tuple[_model.Shape, ...]) -> None: + """Write declarations assigned to one private stub partition.""" + self._write_shapes(shapes) + + def _write_preamble(self) -> None: writer = self.writer writer.line("__all__ = [") with writer.indent(): @@ -212,7 +268,9 @@ def write(self) -> None: with writer.indent(): _write_docstring(writer, _SERVICE_ERROR_DOC) - for shape in self.plan.shapes: + def _write_shapes(self, shapes: Collection[_model.Shape]) -> None: + writer = self.writer + for shape in shapes: writer.line() match shape.type: case _model.ShapeType.STRUCTURE: @@ -223,10 +281,13 @@ def write(self) -> None: self._write_enum(shape) case _: pass + + def _write_unit(self) -> None: + writer = self.writer if self.plan.includes_unit: writer.line() writer.line("@", _DATACLASS, "(kw_only=True)") - writer.line("class Unit(_SerdeShape):") + writer.line("class Unit(", self.serde_shape, "):") with writer.indent(): writer.line("pass") @@ -235,10 +296,17 @@ def _write_structure(self, shape: _model.Shape) -> None: name = self.plan.declaration_names[shape.id] error = shape.has_trait("error") writer.line("@", _DATACLASS, "(kw_only=True)") + bases: tuple[str | _symbols.TypeReference, ...] = ( + (self.service_error, ", ", self.serde_shape) + if error + else (self.serde_shape,) + ) writer.line( "class ", name, - "(ServiceError, _SerdeShape):" if error else "(_SerdeShape):", + "(", + *bases, + "):", ) with writer.indent(): members = tuple( @@ -255,7 +323,7 @@ def _write_structure(self, shape: _model.Shape) -> None: if members or carries_metadata: writer.line() for member in members: - reference = self.symbols.type_reference(member.target) + reference = self._type_reference(member.target) has_default = member.has_trait("default") or self.model.get_target( member ).has_trait("default") @@ -283,20 +351,20 @@ def _write_union(self, shape: _model.Shape) -> None: variants: list[_symbols.TypeReference] = [] for member in shape.members.values(): name = self.plan.union_variant_names[member.id] - variants.append(_symbols.TypeReference(name, self.module)) + variants.append(self._public_reference(name)) writer.line("@", _DATACLASS) - writer.line("class ", name, "(_SerdeShape):") + writer.line("class ", name, "(", self.serde_shape, "):") with writer.indent(): if docs := _documentation(member): _write_docstring(writer, docs) writer.line() - writer.line("value: ", self.symbols.type_reference(member.target)) + writer.line("value: ", self._type_reference(member.target)) writer.line() unknown = self.plan.union_unknown_names[shape.id] - variants.append(_symbols.TypeReference(unknown, self.module)) + variants.append(self._public_reference(unknown)) writer.line("@", _DATACLASS) - writer.line("class ", unknown, "(_SerdeShape):") + writer.line("class ", unknown, "(", self.serde_shape, "):") with writer.indent(): _write_docstring(writer, _UNKNOWN_UNION_DOC) writer.line() @@ -314,7 +382,7 @@ def _write_union(self, shape: _model.Shape) -> None: "def deserialize(self, deserializer: ", _SHAPE_DESERIALIZER, ") -> ", - _symbols.TypeReference(name, self.module), + self._public_reference(name), ": ...", ) @@ -341,12 +409,44 @@ def _write_enum(self, shape: _model.Shape) -> None: if not docs and not shape.members: writer.line("pass") + def _type_reference(self, shape_id: _model.ShapeId) -> _symbols.TypeReference: + return self._partition_reference(self.symbols.type_reference(shape_id)) + + def _public_reference(self, name: str) -> _symbols.TypeReference: + return _symbols.TypeReference( + name, "..models" if self.external_bases else self.module + ) + + def _partition_reference( + self, reference: _symbols.TypeReference + ) -> _symbols.TypeReference: + arguments = tuple(map(self._partition_reference, reference.arguments)) + module = ( + "..models" + if self.external_bases and reference.module == self.module + else reference.module + ) + if module == reference.module and arguments == reference.arguments: + return reference + return replace(reference, module=module, arguments=arguments) + def _model_exports(plan: ModelPlan) -> tuple[str, ...]: exports = ("ServiceError", *plan.registry_index) return (*exports, "Unit") if plan.includes_unit else exports +def _stub_shape_names(plan: ModelPlan, shape: _model.Shape) -> tuple[str, ...]: + declaration = plan.declaration_names[shape.id] + if shape.type is not _model.ShapeType.UNION: + return (declaration,) + return ( + *(plan.union_variant_names[member.id] for member in shape.members.values()), + plan.union_unknown_names[shape.id], + declaration, + ) + + def _write_mapping(writer: PythonWriter, name: str, values: Mapping[str, str]) -> None: if not values: writer.line(name, " = {}") diff --git a/packages/smithy-python/src/smithy_python/schemas.py b/packages/smithy-python/src/smithy_python/schemas.py index c47084e45..cf0764641 100644 --- a/packages/smithy-python/src/smithy_python/schemas.py +++ b/packages/smithy-python/src/smithy_python/schemas.py @@ -27,6 +27,7 @@ class SchemaPlan: """Derived semantic data used to emit compact schemas.""" shapes: tuple[_model.Shape, ...] + partition_groups: tuple[tuple[_model.Shape, ...], ...] names: Mapping[_model.ShapeId, str] strings: Mapping[str, str] traits: Mapping[_model.ShapeId, _TraitPlan] @@ -39,7 +40,7 @@ def from_shapes( symbols: _symbols.SymbolProvider, ) -> SchemaPlan: """Plan schema order, filtered traits, names, and interned strings.""" - emitted = _schema_order(shapes) + emitted, partition_groups = _schema_order(shapes) included = _included_traits(model) inputs = { operation.input for operation in model.iter_shapes(_model.OperationShape) @@ -69,6 +70,7 @@ def from_shapes( return cls( shapes=emitted, + partition_groups=partition_groups, names=MappingProxyType( {shape.id: symbols.schema_name(shape.id) for shape in emitted} ), @@ -83,18 +85,111 @@ class SchemaGenerator: def generate(self, context: GenerationContext[TypesSettings]) -> None: """Register ``._private.schemas`` with the file manifest.""" package = context.settings.package - path = f"src/{package.replace('.', '/')}/_private/schemas.py" - writer = context.files.python_module(path, f"{package}._private.schemas") + directory = package.replace(".", "/") plan = context.plan.schemas - for value, name in plan.strings.items(): + size = context.settings.schema_partition_size + if size is None or len(plan.shapes) <= size: + writer = context.files.python_module( + f"src/{directory}/_private/schemas.py", + f"{package}._private.schemas", + ) + _write_schema_module(writer, plan, plan.shapes) + return + + partitions = _schema_partitions(plan, size) + _write_schema_facade( + context.files.python_module( + f"src/{directory}/_private/schemas.py", + f"{package}._private.schemas", + ), + plan, + partitions, + ) + for index, shapes in enumerate(partitions): + _write_schema_module( + context.files.python_module( + f"src/{directory}/_private/schemas_{index}.py", + f"{package}._private.schemas_{index}", + ), + plan, + shapes, + facade="schemas", + ) + + +_IMPORT_MODULE: Final = _symbols.TypeReference("import_module", "importlib") + + +def _write_schema_facade( + writer: PythonWriter, + plan: SchemaPlan, + partitions: tuple[tuple[_model.Shape, ...], ...], +) -> None: + writer.line("_PARTITION_NAMES: tuple[tuple[str, ...], ...] = (") + with writer.indent(): + for shapes in partitions: + writer.line("(") + with writer.indent(): + for shape in shapes: + writer.line(repr(plan.names[shape.id]), ",") + writer.line("),") + writer.line(")") + writer.line("_PARTITIONS = {") + with writer.indent(): + writer.line("name: index") + writer.line("for index, names in enumerate(_PARTITION_NAMES)") + writer.line("for name in names") + writer.line("}") + writer.line("del _PARTITION_NAMES") + writer.line("def __getattr__(name: str):") + with writer.indent(): + writer.line("try:") + with writer.indent(): + writer.line("partition = _PARTITIONS[name]") + writer.line("except KeyError as error:") + with writer.indent(): + writer.line("raise AttributeError(name) from error") + writer.line( + "module = ", + _IMPORT_MODULE, + '(f".schemas_{partition}", __package__)', + ) + writer.line("value = getattr(module, name)") + writer.line("globals()[name] = value") + writer.line("return value") + writer.line() + writer.line() + writer.line("def __dir__() -> list[str]:") + with writer.indent(): + writer.line("return sorted({*globals(), *_PARTITIONS})") + + +def _write_schema_module( + writer: PythonWriter, + plan: SchemaPlan, + shapes: tuple[_model.Shape, ...], + *, + facade: str | None = None, +) -> None: + local = frozenset(shape.id for shape in shapes) + if facade is not None and any( + member.target in plan.names and member.target not in local + for shape in shapes + for member in shape.members.values() + ): + writer.line("from . import ", facade) + writer.line() + strings = _schema_strings(plan, shapes) + for value, name in plan.strings.items(): + if value in strings: writer.line(name, " = ", repr(value)) - if plan.strings: - writer.line() + if strings: + writer.line() - written: set[_model.ShapeId] = set() - for shape in plan.shapes: - _write_shape(writer, plan, shape, written) - written.add(shape.id) + written: set[_model.ShapeId] = set() + for shape in shapes: + _write_shape(writer, plan, shape, written, local=local, facade=facade) + written.add(shape.id) # This contract is intentionally local to the generator. It must match @@ -264,6 +359,9 @@ def _write_shape( plan: SchemaPlan, shape: _model.Shape, written: set[_model.ShapeId], + *, + local: Collection[_model.ShapeId], + facade: str | None, ) -> None: writer.line(plan.names[shape.id], ": ", _STATIC_SCHEMA, " = (") with writer.indent(): @@ -284,7 +382,13 @@ def _write_shape( for member in shape.members.values(): writer.line( "(", - _schema_ref(plan, member.target, written), + *_schema_ref( + plan, + member.target, + written, + local=local, + facade=facade, + ), ", ", *_trait_parts(plan, plan.traits[member.id]), "),", @@ -406,23 +510,36 @@ def _node_parts( def _schema_ref( - plan: SchemaPlan, target: _model.ShapeId, written: set[_model.ShapeId] -) -> str: + plan: SchemaPlan, + target: _model.ShapeId, + written: set[_model.ShapeId], + *, + local: Collection[_model.ShapeId], + facade: str | None, +) -> tuple[_SourcePart, ...]: if (prelude := _PRELUDE_INTS.get(target)) is not None: - return str(prelude) + return (str(prelude),) try: name = plan.names[target] except KeyError as error: raise _exceptions.CodegenError( f"Schema target {target} was not selected" ) from error - return name if target in written else f"lambda: {name}" + if target in local: + return (name if target in written else f"lambda: {name}",) + if facade is None: + raise _exceptions.CodegenError(f"Schema target {target} is not in this module") + return ("lambda: ", facade, ".", name) def _schema_order( shapes: tuple[_model.Shape, ...], -) -> tuple[_model.Shape, ...]: +) -> tuple[ + tuple[_model.Shape, ...], + tuple[tuple[_model.Shape, ...], ...], +]: by_id = {shape.id: shape for shape in shapes} + positions = {shape.id: index for index, shape in enumerate(shapes)} graph = { shape.id: tuple( member.target for member in shape.members.values() if member.target in by_id @@ -457,6 +574,7 @@ def _schema_order( pending.append((target, 0)) recursive: set[_model.ShapeId] = set() + recursive_components: dict[_model.ShapeId, tuple[_model.ShapeId, ...]] = {} visited.clear() for root in reversed(finished): if root in visited: @@ -473,11 +591,97 @@ def _schema_order( component_pending.append(target) if len(component) > 1 or root in graph[root]: recursive.update(component) + ordered_component = tuple(sorted(component, key=positions.__getitem__)) + for shape_id in component: + recursive_components[shape_id] = ordered_component - return ( + emitted = ( *(by_id[shape_id] for shape_id in finished if shape_id not in recursive), *(shape for shape in shapes if shape.id in recursive), ) + groups: list[tuple[_model.Shape, ...]] = [] + grouped: set[_model.ShapeId] = set() + for shape in emitted: + if shape.id in grouped: + continue + group_ids = recursive_components.get(shape.id, (shape.id,)) + groups.append(tuple(by_id[shape_id] for shape_id in group_ids)) + grouped.update(group_ids) + return emitted, tuple(groups) + + +def _schema_partitions( + plan: SchemaPlan, + size: int, +) -> tuple[tuple[_model.Shape, ...], ...]: + groups = plan.partition_groups + partitions: list[tuple[_model.Shape, ...]] = [] + current: list[_model.Shape] = [] + for group in groups: + if current and len(current) + len(group) > size: + partitions.append(tuple(current)) + current.clear() + current.extend(group) + if current: + partitions.append(tuple(current)) + partition_count = len(partitions) + if partition_count < 2: + return tuple(partitions) + + weights = tuple(_schema_group_weight(plan, group) for group in groups) + partitions.clear() + start = 0 + remaining_weight = sum(weights) + for remaining_parts in range(partition_count, 1, -1): + target = remaining_weight / remaining_parts + end = start + weight = 0 + last_end = len(groups) - remaining_parts + 1 + while end < last_end: + next_weight = weight + weights[end] + if weight and abs(target - weight) <= abs(target - next_weight): + break + weight = next_weight + end += 1 + partitions.append( + tuple(shape for group in groups[start:end] for shape in group) + ) + start = end + remaining_weight -= weight + partitions.append(tuple(shape for group in groups[start:] for shape in group)) + return tuple(partitions) + + +def _schema_group_weight(plan: SchemaPlan, group: tuple[_model.Shape, ...]) -> int: + declarations = sum( + len(plan.strings[value]) + len(repr(value)) + 4 + for value in _schema_strings(plan, group) + ) + return declarations + sum(_schema_shape_weight(plan, shape) for shape in group) + + +def _schema_shape_weight(plan: SchemaPlan, shape: _model.Shape) -> int: + """Estimate compact-schema source bytes without rendering shapes twice.""" + return ( + 60 + + len(shape.id.namespace) + + 2 * len(shape.id.name) + + _source_size(_trait_parts(plan, plan.traits[shape.id])) + + sum( + 40 + + 2 * len(member.name) + + len(plan.names.get(member.target, member.target.name)) + + _source_size(_trait_parts(plan, plan.traits[member.id])) + for member in shape.members.values() + ) + ) + + +def _source_size(parts: Collection[_SourcePart]) -> int: + return sum( + len(part.name if isinstance(part, _symbols.TypeReference) else part) + for part in parts + ) def _intern_traits(strings: _StringStore, traits: _TraitPlan) -> None: @@ -489,6 +693,28 @@ def _intern_traits(strings: _StringStore, traits: _TraitPlan) -> None: strings.intern(value) +def _schema_strings( + plan: SchemaPlan, shapes: Collection[_model.Shape] +) -> frozenset[str]: + values: set[str] = set() + for shape in shapes: + values.update((shape.id.namespace, shape.id.name)) + _collect_trait_strings(values, plan.traits[shape.id]) + for member in shape.members.values(): + values.add(member.name) + _collect_trait_strings(values, plan.traits[member.id]) + return frozenset(values) + + +def _collect_trait_strings(values: set[str], traits: _TraitPlan) -> None: + if isinstance(traits, int): + return + for trait_id, value in traits.items(): + values.add(str(trait_id)) + if isinstance(value, str): + values.add(value) + + class _StringStore: __slots__ = ("_allocated", "_next", "values") diff --git a/packages/smithy-python/tests/unit/test_generation.py b/packages/smithy-python/tests/unit/test_generation.py index 7be37ca29..6065e8b4f 100644 --- a/packages/smithy-python/tests/unit/test_generation.py +++ b/packages/smithy-python/tests/unit/test_generation.py @@ -33,6 +33,15 @@ def test_types_settings_default_to_identity_selection() -> None: ) assert settings.selector == Selector.identity() assert not settings.generate_inputs_and_outputs + assert settings.schema_partition_size == 250 + assert settings.stub_partition_size == 250 + + +def test_types_settings_require_positive_partition_size() -> None: + with pytest.raises(ValueError, match="Partition sizes must be positive"): + TypesSettings("example", "1.0.0", None, schema_partition_size=0) + with pytest.raises(ValueError, match="Partition sizes must be positive"): + TypesSettings("example", "1.0.0", None, stub_partition_size=0) def test_generation_context_shares_manifest_dependencies() -> None: diff --git a/packages/smithy-python/tests/unit/test_models.py b/packages/smithy-python/tests/unit/test_models.py index 2d93ae8d6..db80f53b1 100644 --- a/packages/smithy-python/tests/unit/test_models.py +++ b/packages/smithy-python/tests/unit/test_models.py @@ -44,6 +44,7 @@ def _context( *, service_id: str | None = "example#Service", package: str = _PACKAGE, + stub_partition_size: int | None = 250, ) -> GenerationContext[TypesSettings]: model = load_model(model_data) selection = select_type_shapes( @@ -51,7 +52,12 @@ def _context( service_id=service_id, generate_inputs_and_outputs=True, ) - settings = TypesSettings(package, "1.0.0", None) + settings = TypesSettings( + package, + "1.0.0", + None, + stub_partition_size=stub_partition_size, + ) symbols = SymbolProvider( model, selection, @@ -191,8 +197,15 @@ def _model() -> dict[str, object]: } -def _generate(tmp_path: Path) -> tuple[GenerationContext[TypesSettings], Path]: - context = _context(_model()) +def _generate( + tmp_path: Path, + *, + stub_partition_size: int | None = 250, +) -> tuple[GenerationContext[TypesSettings], Path]: + context = _context( + _model(), + stub_partition_size=stub_partition_size, + ) SchemaGenerator().generate(context) ModelGenerator().generate(context) context.files.text_file(f"src/{_PACKAGE}/__init__.py", "") @@ -433,6 +446,68 @@ def read_struct(self, schema: Any, consumer: Any) -> None: models.Original +def test_partitioned_stub_preserves_public_models( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, source_root = _generate( + tmp_path, + stub_partition_size=2, + ) + package = source_root / _PACKAGE + facade = (package / "models.pyi").read_text() + partitions = sorted((package / "_private").glob("models_*.pyi")) + assert len(partitions) == 4 + assert "class Input(" not in facade + assert "from ._private.models_0 import (" in facade + assert "Input," in facade + assert "Input as Input" not in facade + assert "_SelectionDeserializer" not in facade + assert any("class Input(" in path.read_text() for path in partitions) + + monkeypatch.syspath_prepend( # pyright: ignore[reportUnknownMemberType] + str(source_root) + ) + models = importlib.import_module(f"{_PACKAGE}.models") + assert models.Renamed(value="value").value == "value" + assert models.SelectionText(value="value").value == "value" + + generated = tuple( + str(path) + for path in sorted(package.rglob("*")) + if path.suffix in {".py", ".pyi"} + ) + result = _run_tool("ruff", "check", "--isolated", *generated) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize(("declaration_count", "partition_count"), [(250, 0), (251, 2)]) +def test_stub_partition_threshold( + tmp_path: Path, + declaration_count: int, + partition_count: int, +) -> None: + package = "generated_threshold" + context = _context( + { + "smithy": "2.0", + "shapes": { + f"example#Shape{index}": {"type": "structure"} + for index in range(declaration_count) + }, + }, + service_id=None, + package=package, + ) + ModelGenerator().generate(context) + context.files.commit(tmp_path) + + source = tmp_path / f"src/{package}" + assert len(tuple((source / "_private").glob("models_*.pyi"))) == partition_count + facade = (source / "models.pyi").read_text() + assert ("._private.models_0" in facade) is bool(partition_count) + + def test_generated_model_names_do_not_shadow_runtime_helpers( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -750,10 +825,12 @@ def test_modeled_unit_does_not_shadow_prelude_unit( assert isinstance(decoded.action.value, models.Unit) +@pytest.mark.parametrize("stub_partition_size", [None, 2]) def test_generated_stub_passes_strict_pyright( tmp_path: Path, + stub_partition_size: int | None, ) -> None: - _, source_root = _generate(tmp_path) + _, source_root = _generate(tmp_path, stub_partition_size=stub_partition_size) (tmp_path / "check.py").write_text( ( f"from {_PACKAGE}.models import Input, Problem, Renamed, SelectionText, Tone\n\n" diff --git a/packages/smithy-python/tests/unit/test_schemas.py b/packages/smithy-python/tests/unit/test_schemas.py index ca5bc72cb..3a8b81f06 100644 --- a/packages/smithy-python/tests/unit/test_schemas.py +++ b/packages/smithy-python/tests/unit/test_schemas.py @@ -4,7 +4,9 @@ from __future__ import annotations +import importlib import runpy +import sys from collections.abc import Mapping from pathlib import Path from types import MappingProxyType @@ -32,10 +34,19 @@ from smithy_python.symbols import SymbolProvider -def _context(model_data: Mapping[str, object]) -> GenerationContext[TypesSettings]: +def _context( + model_data: Mapping[str, object], + *, + schema_partition_size: int | None = 250, +) -> GenerationContext[TypesSettings]: model = load_model(model_data) selection = select_type_shapes(model, generate_inputs_and_outputs=True) - settings = TypesSettings("example.types", "1.0.0", None) + settings = TypesSettings( + "example.types", + "1.0.0", + None, + schema_partition_size=schema_partition_size, + ) symbols = SymbolProvider(model, selection, package=settings.package) dependencies = DependencySet() return GenerationContext( @@ -231,6 +242,158 @@ def test_schema_generator_emits_lazy_hydratable_schemas(tmp_path: Path) -> None: assert str(node.members_by_index[2].expect_member_target().id) == "example#Choice" +def test_schema_generator_partitions_behind_lazy_facade( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + context = _context( + { + "smithy": "2.0", + "shapes": { + "example#Leaf": {"type": "string"}, + "example#Node": { + "type": "structure", + "members": {"leaf": {"target": "example#Leaf"}}, + }, + "example#CycleA": { + "type": "structure", + "members": {"b": {"target": "example#CycleB"}}, + }, + "example#CycleB": { + "type": "structure", + "members": {"a": {"target": "example#CycleA"}}, + }, + }, + }, + schema_partition_size=1, + ) + SchemaGenerator().generate(context) + context.files.text_file("src/example/types/__init__.py", "") + context.files.text_file("src/example/types/_private/__init__.py", "") + context.files.commit(tmp_path) + + package = tmp_path / "src/example/types/_private" + assert (package / "schemas.py").is_file() + assert (package / "schemas_0.py").is_file() + assert (package / "schemas_1.py").is_file() + facade_source = (package / "schemas.py").read_text() + assert "_PARTITION_NAMES: tuple[tuple[str, ...], ...]" in facade_source + assert "name: index" in facade_source + assert facade_source.count(".schemas_") == 1 + partitions = tuple(package.glob("schemas_*.py")) + cycle = next( + path.read_text() for path in partitions if "CYCLE_A:" in path.read_text() + ) + assert "CYCLE_B:" in cycle + + monkeypatch.syspath_prepend( # pyright: ignore[reportUnknownMemberType] + str(tmp_path / "src") + ) + facade = importlib.import_module("example.types._private.schemas") + assert "example.types._private.schemas_0" not in sys.modules + assert "example.types._private.schemas_1" not in sys.modules + + static_node = cast(tuple[Any, ...], facade.NODE) + assert facade.NODE is static_node + assert "NODE" in dir(facade) + assert "example.types._private.schemas_1" in sys.modules + assert "example.types._private.schemas_0" not in sys.modules + with pytest.raises(AttributeError): + facade.MISSING + node = _schema_compact.hydrate(static_node) + assert str(node.members["leaf"].expect_member_target().id) == "example#Leaf" + assert "example.types._private.schemas_0" in sys.modules + cycle_a = _schema_compact.hydrate(cast(tuple[Any, ...], facade.CYCLE_A)) + cycle_b = cycle_a.members["b"].expect_member_target() + assert cycle_b.members["a"].expect_member_target() is cycle_a + + +@pytest.mark.parametrize(("shape_count", "partition_count"), [(250, 0), (251, 2)]) +def test_schema_partition_threshold( + tmp_path: Path, + shape_count: int, + partition_count: int, +) -> None: + context = _context( + { + "smithy": "2.0", + "shapes": { + f"example#Shape{index}": {"type": "structure"} + for index in range(shape_count) + }, + }, + ) + SchemaGenerator().generate(context) + context.files.commit(tmp_path) + + private = tmp_path / "src/example/types/_private" + assert len(tuple(private.glob("schemas_*.py"))) == partition_count + facade = (private / "schemas.py").read_text() + assert ("_PARTITIONS" in facade) is bool(partition_count) + + +def test_schema_partitions_balance_emitted_size(tmp_path: Path) -> None: + context = _context( + { + "smithy": "2.0", + "shapes": { + "example#AHeavy": { + "type": "structure", + "members": { + f"member{index}": {"target": "smithy.api#String"} + for index in range(20) + }, + }, + "example#BTiny": {"type": "structure"}, + "example#CTiny": {"type": "structure"}, + "example#DTiny": {"type": "structure"}, + }, + }, + schema_partition_size=2, + ) + SchemaGenerator().generate(context) + context.files.commit(tmp_path) + private = tmp_path / "src/example/types/_private" + + first = (private / "schemas_0.py").read_text() + second = (private / "schemas_1.py").read_text() + assert "A_HEAVY:" in first + assert "B_TINY:" not in first + assert all(f"{name}:" in second for name in ("B_TINY", "C_TINY", "D_TINY")) + + +def test_schema_partition_weights_include_interned_strings(tmp_path: Path) -> None: + context = _context( + { + "smithy": "2.0", + "shapes": { + "example#AHeavy": { + "type": "structure", + "members": { + "value": { + "target": "smithy.api#String", + "traits": {"smithy.api#default": "x" * 10_000}, + } + }, + }, + "example#BTiny": {"type": "structure"}, + "example#CTiny": {"type": "structure"}, + "example#DTiny": {"type": "structure"}, + }, + }, + schema_partition_size=2, + ) + SchemaGenerator().generate(context) + context.files.commit(tmp_path) + private = tmp_path / "src/example/types/_private" + + first = (private / "schemas_0.py").read_text() + second = (private / "schemas_1.py").read_text() + assert "A_HEAVY:" in first + assert "B_TINY:" not in first + assert all(f"{name}:" in second for name in ("B_TINY", "C_TINY", "D_TINY")) + + def test_schema_generator_preserves_primitive_prelude_targets(tmp_path: Path) -> None: context = _context( {