Add size hits to list and map serializers

codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/generators/ListGenerator.java

@@ -45,7 +45,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:

codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/generators/MapGenerator.java

@@ -47,7 +47,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}

python-packages/smithy-core/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

python-packages/smithy-core/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,23 +245,23 @@ 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
         else:
             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:

python-packages/smithy-core/tests/unit/test_documents.py

@@ -623,21 +623,21 @@ 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:
             serializer.write_struct(SCHEMA.members["structMember"], self.struct_member)
         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:

python-packages/smithy-json/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:

python-packages/smithy-json/tests/unit/__init__.py

@@ -170,21 +170,21 @@ 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:
             serializer.write_struct(SCHEMA.members["structMember"], self.struct_member)
         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))

python-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: