Skip to content

Commit 68252bd

Browse files
jonathan343jonathan343
authored
Merge pull request #382 from jonathan343/add-size-hints
Add size hits to list and map serializers
2 parents d9c4928 + c204e9b commit 68252bd

File tree

8 files changed

+58
-30
lines changed

8 files changed

+58
-30
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ private void generateSerializer() {
4444
writer.write("""
4545
def $1L(serializer: ShapeSerializer, schema: Schema, value: $2T) -> None:
4646
member_schema = schema.members["member"]
47-
with serializer.begin_list(schema) as ls:
47+
with serializer.begin_list(schema, len(value)) as ls:
4848
for e in value:
4949
${?sparse}
5050
if e is None:

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ private void generateSerializer() {
4646
// narrowing out the None even though there's an explicit is None check.
4747
writer.write("""
4848
def $1L(serializer: ShapeSerializer, schema: Schema, value: $2T) -> None:
49-
with serializer.begin_map(schema) as m:
49+
with serializer.begin_map(schema, len(value)) as m:
5050
value_schema = schema.members["value"]
5151
for k, v in value.items():
5252
${?sparse}

‎packages/smithy-core/src/smithy_core/documents.py‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -268,12 +268,18 @@ def serialize_contents(self, serializer: ShapeSerializer) -> None:
268268
with serializer.begin_struct(self._schema) as struct_serializer:
269269
self.serialize_members(struct_serializer)
270270
case ShapeType.LIST:
271-
with serializer.begin_list(self._schema) as list_serializer:
272-
for element in self.as_list():
271+
list_value = self.as_list()
272+
with serializer.begin_list(
273+
self._schema, len(list_value)
274+
) as list_serializer:
275+
for element in list_value:
273276
element.serialize(list_serializer)
274277
case ShapeType.MAP:
275-
with serializer.begin_map(self._schema) as map_serializer:
276-
for key, value in self.as_map().items():
278+
map_value = self.as_map()
279+
with serializer.begin_map(
280+
self._schema, len(map_value)
281+
) as map_serializer:
282+
for key, value in map_value.items():
277283
map_serializer.entry(key, lambda s: value.serialize(s))
278284
case ShapeType.STRING | ShapeType.ENUM:
279285
serializer.write_string(self._schema, self.as_string())
@@ -427,7 +433,7 @@ def begin_struct(self, schema: "Schema") -> Iterator[ShapeSerializer]:
427433

428434
@override
429435
@contextmanager
430-
def begin_list(self, schema: "Schema") -> Iterator[ShapeSerializer]:
436+
def begin_list(self, schema: "Schema", size: int) -> Iterator[ShapeSerializer]:
431437
delegate = _DocumentListSerializer(schema)
432438
try:
433439
yield delegate
@@ -438,7 +444,7 @@ def begin_list(self, schema: "Schema") -> Iterator[ShapeSerializer]:
438444

439445
@override
440446
@contextmanager
441-
def begin_map(self, schema: "Schema") -> Iterator[MapSerializer]:
447+
def begin_map(self, schema: "Schema", size: int) -> Iterator[MapSerializer]:
442448
delegate = _DocumentMapSerializer(schema)
443449
try:
444450
yield delegate

‎packages/smithy-core/src/smithy_core/serializers.py‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,9 @@ def write_struct(self, schema: "Schema", struct: "SerializeableStruct") -> None:
5252
with self.begin_struct(schema=schema) as struct_serializer:
5353
struct.serialize_members(struct_serializer)
5454

55-
def begin_list(self, schema: "Schema") -> AbstractContextManager["ShapeSerializer"]:
55+
def begin_list(
56+
self, schema: "Schema", size: int
57+
) -> AbstractContextManager["ShapeSerializer"]:
5658
"""Open a list for writing.
5759
5860
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
6264
inserting any data needed between elements.
6365
6466
:param schema: The schema of the list.
67+
:param size: The size of the list.
6568
:returns: A context manager containing an element serializer.
6669
"""
6770
...
6871

69-
def begin_map(self, schema: "Schema") -> AbstractContextManager["MapSerializer"]:
72+
def begin_map(
73+
self, schema: "Schema", size: int
74+
) -> AbstractContextManager["MapSerializer"]:
7075
"""Open a map for writing.
7176
7277
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"]
7782
data needed between entries.
7883
7984
:param schema: The schema of the map.
85+
:param size: The size of the map.
8086
:returns: A context manager containing a map serializer.
8187
"""
8288
...
@@ -239,23 +245,23 @@ def begin_struct(self, schema: "Schema") -> Iterator[ShapeSerializer]:
239245
self.after(schema)
240246

241247
@contextmanager
242-
def begin_list(self, schema: "Schema") -> Iterator[ShapeSerializer]:
248+
def begin_list(self, schema: "Schema", size: int) -> Iterator[ShapeSerializer]:
243249
delegate = self.before(schema)
244250

245251
try:
246-
with delegate.begin_list(schema) as s:
252+
with delegate.begin_list(schema, size) as s:
247253
yield s
248254
except Exception:
249255
raise
250256
else:
251257
self.after(schema)
252258

253259
@contextmanager
254-
def begin_map(self, schema: "Schema") -> Iterator[MapSerializer]:
260+
def begin_map(self, schema: "Schema", size: int) -> Iterator[MapSerializer]:
255261
delegate = self.before(schema)
256262

257263
try:
258-
with delegate.begin_map(schema) as s:
264+
with delegate.begin_map(schema, size) as s:
259265
yield s
260266
except Exception:
261267
raise
@@ -335,10 +341,14 @@ def begin_struct(
335341
) -> AbstractContextManager["ShapeSerializer"]:
336342
self._invalid_state(schema)
337343

338-
def begin_list(self, schema: "Schema") -> AbstractContextManager["ShapeSerializer"]:
344+
def begin_list(
345+
self, schema: "Schema", size: int
346+
) -> AbstractContextManager["ShapeSerializer"]:
339347
self._invalid_state(schema)
340348

341-
def begin_map(self, schema: "Schema") -> AbstractContextManager["MapSerializer"]:
349+
def begin_map(
350+
self, schema: "Schema", size: int
351+
) -> AbstractContextManager["MapSerializer"]:
342352
self._invalid_state(schema)
343353

344354
def write_null(self, schema: "Schema") -> None:

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

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -623,21 +623,21 @@ def serialize_members(self, serializer: ShapeSerializer) -> None:
623623
if self.list_member is not None:
624624
schema = SCHEMA.members["listMember"]
625625
target_schema = schema.expect_member_target().members["member"]
626-
with serializer.begin_list(schema) as ls:
626+
with serializer.begin_list(schema, len(self.list_member)) as ls:
627627
for element in self.list_member:
628628
ls.write_string(target_schema, element)
629629
if self.map_member is not None:
630630
schema = SCHEMA.members["mapMember"]
631631
target_schema = schema.expect_member_target().members["value"]
632-
with serializer.begin_map(schema) as ms:
632+
with serializer.begin_map(schema, len(self.map_member)) as ms:
633633
for key, value in self.map_member.items():
634634
ms.entry(key, lambda vs: vs.write_string(target_schema, value)) # type: ignore
635635
if self.struct_member is not None:
636636
serializer.write_struct(SCHEMA.members["structMember"], self.struct_member)
637637
if self.sparse_list_member is not None:
638638
schema = SCHEMA.members["sparseListMember"]
639639
target_schema = schema.expect_member_target().members["member"]
640-
with serializer.begin_list(schema) as ls:
640+
with serializer.begin_list(schema, len(self.sparse_list_member)) as ls:
641641
for element in self.sparse_list_member:
642642
if element is None:
643643
ls.write_null(target_schema)
@@ -646,7 +646,7 @@ def serialize_members(self, serializer: ShapeSerializer) -> None:
646646
if self.sparse_map_member is not None:
647647
schema = SCHEMA.members["sparseMapMember"]
648648
target_schema = schema.expect_member_target().members["value"]
649-
with serializer.begin_map(schema) as ms:
649+
with serializer.begin_map(schema, len(self.sparse_map_member)) as ms:
650650
for key, value in self.sparse_map_member.items():
651651
if value is None:
652652
ms.entry(key, lambda vs: vs.write_null(target_schema))
@@ -840,7 +840,9 @@ def test_document_serializer(given: Any, expected: Document):
840840
serializer.write_document(DOCUMENT, given)
841841
case list():
842842
given = cast(list[str], given)
843-
with serializer.begin_list(SPARSE_STRING_LIST_SCHEMA) as list_serializer:
843+
with serializer.begin_list(
844+
SPARSE_STRING_LIST_SCHEMA, len(given)
845+
) as list_serializer:
844846
member_schema = SPARSE_STRING_LIST_SCHEMA.members["member"]
845847
for e in given:
846848
if e is None:
@@ -849,7 +851,9 @@ def test_document_serializer(given: Any, expected: Document):
849851
list_serializer.write_string(member_schema, e)
850852
case dict():
851853
given = cast(dict[str, str], given)
852-
with serializer.begin_map(SPARSE_STRING_MAP_SCHEMA) as map_serializer:
854+
with serializer.begin_map(
855+
SPARSE_STRING_MAP_SCHEMA, len(given)
856+
) as map_serializer:
853857
member_schema = SPARSE_STRING_MAP_SCHEMA.members["value"]
854858
for k, v in given.items():
855859
if v is None:

‎packages/smithy-json/src/smithy_json/_private/serializers.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,10 +50,14 @@ def begin_struct(
5050
) -> AbstractContextManager["ShapeSerializer"]:
5151
return JSONStructSerializer(self._stream, self, self._use_json_name)
5252

53-
def begin_list(self, schema: "Schema") -> AbstractContextManager["ShapeSerializer"]:
53+
def begin_list(
54+
self, schema: "Schema", size: int
55+
) -> AbstractContextManager["ShapeSerializer"]:
5456
return JSONListSerializer(self._stream, self)
5557

56-
def begin_map(self, schema: "Schema") -> AbstractContextManager["MapSerializer"]:
58+
def begin_map(
59+
self, schema: "Schema", size: int
60+
) -> AbstractContextManager["MapSerializer"]:
5761
return JSONMapSerializer(self._stream, self)
5862

5963
def write_null(self, schema: "Schema") -> None:

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -170,21 +170,21 @@ def serialize_members(self, serializer: ShapeSerializer) -> None:
170170
if self.list_member is not None:
171171
schema = SCHEMA.members["listMember"]
172172
target_schema = schema.expect_member_target().members["member"]
173-
with serializer.begin_list(schema) as ls:
173+
with serializer.begin_list(schema, len(self.list_member)) as ls:
174174
for element in self.list_member:
175175
ls.write_string(target_schema, element)
176176
if self.map_member is not None:
177177
schema = SCHEMA.members["mapMember"]
178178
target_schema = schema.expect_member_target().members["value"]
179-
with serializer.begin_map(schema) as ms:
179+
with serializer.begin_map(schema, len(self.map_member)) as ms:
180180
for key, value in self.map_member.items():
181181
ms.entry(key, lambda vs: vs.write_string(target_schema, value)) # type: ignore
182182
if self.struct_member is not None:
183183
serializer.write_struct(SCHEMA.members["structMember"], self.struct_member)
184184
if self.sparse_list_member is not None:
185185
schema = SCHEMA.members["sparseListMember"]
186186
target_schema = schema.expect_member_target().members["member"]
187-
with serializer.begin_list(schema) as ls:
187+
with serializer.begin_list(schema, len(self.sparse_list_member)) as ls:
188188
for element in self.sparse_list_member:
189189
if element is None:
190190
ls.write_null(target_schema)
@@ -193,7 +193,7 @@ def serialize_members(self, serializer: ShapeSerializer) -> None:
193193
if self.sparse_map_member is not None:
194194
schema = SCHEMA.members["sparseMapMember"]
195195
target_schema = schema.expect_member_target().members["value"]
196-
with serializer.begin_map(schema) as ms:
196+
with serializer.begin_map(schema, len(self.sparse_map_member)) as ms:
197197
for key, value in self.sparse_map_member.items():
198198
if value is None:
199199
ms.entry(key, lambda vs: vs.write_null(target_schema))

‎packages/smithy-json/tests/unit/test_serializers.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,9 @@ def test_json_serializer(given: Any, expected: bytes) -> None:
5050
serializer.write_document(given._schema, given) # type: ignore
5151
case list():
5252
given = cast(list[str], given)
53-
with serializer.begin_list(SPARSE_STRING_LIST_SCHEMA) as list_serializer:
53+
with serializer.begin_list(
54+
SPARSE_STRING_LIST_SCHEMA, len(given)
55+
) as list_serializer:
5456
member_schema = SPARSE_STRING_LIST_SCHEMA.members["member"]
5557
for e in given:
5658
if e is None:
@@ -59,7 +61,9 @@ def test_json_serializer(given: Any, expected: bytes) -> None:
5961
list_serializer.write_string(member_schema, e)
6062
case dict():
6163
given = cast(dict[str, str], given)
62-
with serializer.begin_map(SPARSE_STRING_MAP_SCHEMA) as map_serializer:
64+
with serializer.begin_map(
65+
SPARSE_STRING_MAP_SCHEMA, len(given)
66+
) as map_serializer:
6367
member_schema = SPARSE_STRING_MAP_SCHEMA.members["value"]
6468
for k, v in given.items():
6569
if v is None:

0 commit comments

Comments
 (0)