diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ListGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ListGenerator.java index 91f60c38f..7fe3c2d6e 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ListGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/ListGenerator.java @@ -44,7 +44,7 @@ private void generateSerializer() { writer.write(""" def $1L(serializer: ShapeSerializer, schema: Schema, value: $2T) -> None: member_schema = schema.members["member"] - with serializer.begin_list(schema) as ls: + with serializer.begin_list(schema, len(value)) as ls: for e in value: ${?sparse} if e is None: diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/MapGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/MapGenerator.java index 853ee722a..0923d7ee7 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/MapGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/MapGenerator.java @@ -46,7 +46,7 @@ private void generateSerializer() { // narrowing out the None even though there's an explicit is None check. writer.write(""" def $1L(serializer: ShapeSerializer, schema: Schema, value: $2T) -> None: - with serializer.begin_map(schema) as m: + with serializer.begin_map(schema, len(value)) as m: value_schema = schema.members["value"] for k, v in value.items(): ${?sparse} diff --git a/packages/smithy-core/src/smithy_core/documents.py b/packages/smithy-core/src/smithy_core/documents.py index 1a29d2128..89328a7a7 100644 --- a/packages/smithy-core/src/smithy_core/documents.py +++ b/packages/smithy-core/src/smithy_core/documents.py @@ -268,12 +268,18 @@ def serialize_contents(self, serializer: ShapeSerializer) -> None: with serializer.begin_struct(self._schema) as struct_serializer: self.serialize_members(struct_serializer) case ShapeType.LIST: - with serializer.begin_list(self._schema) as list_serializer: - for element in self.as_list(): + list_value = self.as_list() + with serializer.begin_list( + self._schema, len(list_value) + ) as list_serializer: + for element in list_value: element.serialize(list_serializer) case ShapeType.MAP: - with serializer.begin_map(self._schema) as map_serializer: - for key, value in self.as_map().items(): + map_value = self.as_map() + with serializer.begin_map( + self._schema, len(map_value) + ) as map_serializer: + for key, value in map_value.items(): map_serializer.entry(key, lambda s: value.serialize(s)) case ShapeType.STRING | ShapeType.ENUM: serializer.write_string(self._schema, self.as_string()) @@ -427,7 +433,7 @@ def begin_struct(self, schema: "Schema") -> Iterator[ShapeSerializer]: @override @contextmanager - def begin_list(self, schema: "Schema") -> Iterator[ShapeSerializer]: + def begin_list(self, schema: "Schema", size: int) -> Iterator[ShapeSerializer]: delegate = _DocumentListSerializer(schema) try: yield delegate @@ -438,7 +444,7 @@ def begin_list(self, schema: "Schema") -> Iterator[ShapeSerializer]: @override @contextmanager - def begin_map(self, schema: "Schema") -> Iterator[MapSerializer]: + def begin_map(self, schema: "Schema", size: int) -> Iterator[MapSerializer]: delegate = _DocumentMapSerializer(schema) try: yield delegate diff --git a/packages/smithy-core/src/smithy_core/serializers.py b/packages/smithy-core/src/smithy_core/serializers.py index 6a4648927..2da8eab99 100644 --- a/packages/smithy-core/src/smithy_core/serializers.py +++ b/packages/smithy-core/src/smithy_core/serializers.py @@ -52,7 +52,9 @@ def write_struct(self, schema: "Schema", struct: "SerializeableStruct") -> None: with self.begin_struct(schema=schema) as struct_serializer: struct.serialize_members(struct_serializer) - def begin_list(self, schema: "Schema") -> AbstractContextManager["ShapeSerializer"]: + def begin_list( + self, schema: "Schema", size: int + ) -> AbstractContextManager["ShapeSerializer"]: """Open a list for writing. The returned context manager is responsible for closing the list when the caller @@ -62,11 +64,14 @@ def begin_list(self, schema: "Schema") -> AbstractContextManager["ShapeSerialize inserting any data needed between elements. :param schema: The schema of the list. + :param size: The size of the list. :returns: A context manager containing an element serializer. """ ... - def begin_map(self, schema: "Schema") -> AbstractContextManager["MapSerializer"]: + def begin_map( + self, schema: "Schema", size: int + ) -> AbstractContextManager["MapSerializer"]: """Open a map for writing. The returned context manager is responsible for closing the map when the caller @@ -77,6 +82,7 @@ def begin_map(self, schema: "Schema") -> AbstractContextManager["MapSerializer"] data needed between entries. :param schema: The schema of the map. + :param size: The size of the map. :returns: A context manager containing a map serializer. """ ... @@ -239,11 +245,11 @@ def begin_struct(self, schema: "Schema") -> Iterator[ShapeSerializer]: self.after(schema) @contextmanager - def begin_list(self, schema: "Schema") -> Iterator[ShapeSerializer]: + def begin_list(self, schema: "Schema", size: int) -> Iterator[ShapeSerializer]: delegate = self.before(schema) try: - with delegate.begin_list(schema) as s: + with delegate.begin_list(schema, size) as s: yield s except Exception: raise @@ -251,11 +257,11 @@ def begin_list(self, schema: "Schema") -> Iterator[ShapeSerializer]: self.after(schema) @contextmanager - def begin_map(self, schema: "Schema") -> Iterator[MapSerializer]: + def begin_map(self, schema: "Schema", size: int) -> Iterator[MapSerializer]: delegate = self.before(schema) try: - with delegate.begin_map(schema) as s: + with delegate.begin_map(schema, size) as s: yield s except Exception: raise @@ -335,10 +341,14 @@ def begin_struct( ) -> AbstractContextManager["ShapeSerializer"]: self._invalid_state(schema) - def begin_list(self, schema: "Schema") -> AbstractContextManager["ShapeSerializer"]: + def begin_list( + self, schema: "Schema", size: int + ) -> AbstractContextManager["ShapeSerializer"]: self._invalid_state(schema) - def begin_map(self, schema: "Schema") -> AbstractContextManager["MapSerializer"]: + def begin_map( + self, schema: "Schema", size: int + ) -> AbstractContextManager["MapSerializer"]: self._invalid_state(schema) def write_null(self, schema: "Schema") -> None: diff --git a/packages/smithy-core/tests/unit/test_documents.py b/packages/smithy-core/tests/unit/test_documents.py index 3afea0fe9..4f7fd93b0 100644 --- a/packages/smithy-core/tests/unit/test_documents.py +++ b/packages/smithy-core/tests/unit/test_documents.py @@ -623,13 +623,13 @@ def serialize_members(self, serializer: ShapeSerializer) -> None: if self.list_member is not None: schema = SCHEMA.members["listMember"] target_schema = schema.expect_member_target().members["member"] - with serializer.begin_list(schema) as ls: + with serializer.begin_list(schema, len(self.list_member)) as ls: for element in self.list_member: ls.write_string(target_schema, element) if self.map_member is not None: schema = SCHEMA.members["mapMember"] target_schema = schema.expect_member_target().members["value"] - with serializer.begin_map(schema) as ms: + with serializer.begin_map(schema, len(self.map_member)) as ms: for key, value in self.map_member.items(): ms.entry(key, lambda vs: vs.write_string(target_schema, value)) # type: ignore if self.struct_member is not None: @@ -637,7 +637,7 @@ def serialize_members(self, serializer: ShapeSerializer) -> None: if self.sparse_list_member is not None: schema = SCHEMA.members["sparseListMember"] target_schema = schema.expect_member_target().members["member"] - with serializer.begin_list(schema) as ls: + with serializer.begin_list(schema, len(self.sparse_list_member)) as ls: for element in self.sparse_list_member: if element is None: ls.write_null(target_schema) @@ -646,7 +646,7 @@ def serialize_members(self, serializer: ShapeSerializer) -> None: if self.sparse_map_member is not None: schema = SCHEMA.members["sparseMapMember"] target_schema = schema.expect_member_target().members["value"] - with serializer.begin_map(schema) as ms: + with serializer.begin_map(schema, len(self.sparse_map_member)) as ms: for key, value in self.sparse_map_member.items(): if value is None: ms.entry(key, lambda vs: vs.write_null(target_schema)) @@ -840,7 +840,9 @@ def test_document_serializer(given: Any, expected: Document): serializer.write_document(DOCUMENT, given) case list(): given = cast(list[str], given) - with serializer.begin_list(SPARSE_STRING_LIST_SCHEMA) as list_serializer: + with serializer.begin_list( + SPARSE_STRING_LIST_SCHEMA, len(given) + ) as list_serializer: member_schema = SPARSE_STRING_LIST_SCHEMA.members["member"] for e in given: if e is None: @@ -849,7 +851,9 @@ def test_document_serializer(given: Any, expected: Document): list_serializer.write_string(member_schema, e) case dict(): given = cast(dict[str, str], given) - with serializer.begin_map(SPARSE_STRING_MAP_SCHEMA) as map_serializer: + with serializer.begin_map( + SPARSE_STRING_MAP_SCHEMA, len(given) + ) as map_serializer: member_schema = SPARSE_STRING_MAP_SCHEMA.members["value"] for k, v in given.items(): if v is None: diff --git a/packages/smithy-json/src/smithy_json/_private/serializers.py b/packages/smithy-json/src/smithy_json/_private/serializers.py index 217ca67ad..91f8eb50f 100644 --- a/packages/smithy-json/src/smithy_json/_private/serializers.py +++ b/packages/smithy-json/src/smithy_json/_private/serializers.py @@ -50,10 +50,14 @@ def begin_struct( ) -> AbstractContextManager["ShapeSerializer"]: return JSONStructSerializer(self._stream, self, self._use_json_name) - def begin_list(self, schema: "Schema") -> AbstractContextManager["ShapeSerializer"]: + def begin_list( + self, schema: "Schema", size: int + ) -> AbstractContextManager["ShapeSerializer"]: return JSONListSerializer(self._stream, self) - def begin_map(self, schema: "Schema") -> AbstractContextManager["MapSerializer"]: + def begin_map( + self, schema: "Schema", size: int + ) -> AbstractContextManager["MapSerializer"]: return JSONMapSerializer(self._stream, self) def write_null(self, schema: "Schema") -> None: diff --git a/packages/smithy-json/tests/unit/__init__.py b/packages/smithy-json/tests/unit/__init__.py index c61d173c6..d82feb010 100644 --- a/packages/smithy-json/tests/unit/__init__.py +++ b/packages/smithy-json/tests/unit/__init__.py @@ -170,13 +170,13 @@ def serialize_members(self, serializer: ShapeSerializer) -> None: if self.list_member is not None: schema = SCHEMA.members["listMember"] target_schema = schema.expect_member_target().members["member"] - with serializer.begin_list(schema) as ls: + with serializer.begin_list(schema, len(self.list_member)) as ls: for element in self.list_member: ls.write_string(target_schema, element) if self.map_member is not None: schema = SCHEMA.members["mapMember"] target_schema = schema.expect_member_target().members["value"] - with serializer.begin_map(schema) as ms: + with serializer.begin_map(schema, len(self.map_member)) as ms: for key, value in self.map_member.items(): ms.entry(key, lambda vs: vs.write_string(target_schema, value)) # type: ignore if self.struct_member is not None: @@ -184,7 +184,7 @@ def serialize_members(self, serializer: ShapeSerializer) -> None: if self.sparse_list_member is not None: schema = SCHEMA.members["sparseListMember"] target_schema = schema.expect_member_target().members["member"] - with serializer.begin_list(schema) as ls: + with serializer.begin_list(schema, len(self.sparse_list_member)) as ls: for element in self.sparse_list_member: if element is None: ls.write_null(target_schema) @@ -193,7 +193,7 @@ def serialize_members(self, serializer: ShapeSerializer) -> None: if self.sparse_map_member is not None: schema = SCHEMA.members["sparseMapMember"] target_schema = schema.expect_member_target().members["value"] - with serializer.begin_map(schema) as ms: + with serializer.begin_map(schema, len(self.sparse_map_member)) as ms: for key, value in self.sparse_map_member.items(): if value is None: ms.entry(key, lambda vs: vs.write_null(target_schema)) diff --git a/packages/smithy-json/tests/unit/test_serializers.py b/packages/smithy-json/tests/unit/test_serializers.py index 567474687..82b525f2e 100644 --- a/packages/smithy-json/tests/unit/test_serializers.py +++ b/packages/smithy-json/tests/unit/test_serializers.py @@ -50,7 +50,9 @@ def test_json_serializer(given: Any, expected: bytes) -> None: serializer.write_document(given._schema, given) # type: ignore case list(): given = cast(list[str], given) - with serializer.begin_list(SPARSE_STRING_LIST_SCHEMA) as list_serializer: + with serializer.begin_list( + SPARSE_STRING_LIST_SCHEMA, len(given) + ) as list_serializer: member_schema = SPARSE_STRING_LIST_SCHEMA.members["member"] for e in given: if e is None: @@ -59,7 +61,9 @@ def test_json_serializer(given: Any, expected: bytes) -> None: list_serializer.write_string(member_schema, e) case dict(): given = cast(dict[str, str], given) - with serializer.begin_map(SPARSE_STRING_MAP_SCHEMA) as map_serializer: + with serializer.begin_map( + SPARSE_STRING_MAP_SCHEMA, len(given) + ) as map_serializer: member_schema = SPARSE_STRING_MAP_SCHEMA.members["value"] for k, v in given.items(): if v is None: