Add size hits to list and map serializers

Author: jonathan343

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/aio/test_types.py

@@ -5,7 +5,6 @@
 from typing import Self
 
 import pytest
-
 from smithy_core.aio.types import (
     AsyncBytesProvider,
     AsyncBytesReader,

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

@@ -5,7 +5,6 @@
 from typing import Any, Self, cast
 
 import pytest
-
 from smithy_core.deserializers import ShapeDeserializer
 from smithy_core.documents import (
     Document,
@@ -623,21 +622,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 +645,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 +839,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 +850,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-core/tests/unit/test_identity.py

@@ -4,7 +4,6 @@
 
 import pytest
 from freezegun import freeze_time
-
 from smithy_core.identity import Identity
 
 

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

@@ -2,7 +2,6 @@
 #  SPDX-License-Identifier: Apache-2.0
 
 import pytest
-
 from smithy_core.exceptions import SmithyRetryException
 from smithy_core.interfaces.retries import RetryErrorInfo, RetryErrorType
 from smithy_core.retries import ExponentialBackoffJitterType as EBJT

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

@@ -1,7 +1,6 @@
 from dataclasses import replace
 
 import pytest
-
 from smithy_core.exceptions import ExpectationNotMetException
 from smithy_core.schemas import Schema
 from smithy_core.shapes import ShapeID, ShapeType

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

@@ -1,5 +1,4 @@
 import pytest
-
 from smithy_core.exceptions import ExpectationNotMetException, SmithyException
 from smithy_core.shapes import ShapeID
 

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

@@ -5,7 +5,6 @@
 from datetime import datetime, timezone
 
 import pytest
-
 from smithy_core.exceptions import ExpectationNotMetException
 from smithy_core.types import JsonBlob, JsonString, TimestampFormat
 

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

@@ -1,7 +1,6 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
 import pytest
-
 from smithy_core import URI, HostType
 from smithy_core.exceptions import SmithyException