Generate schema-based deserializers

codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/SetupGenerator.java

@@ -111,6 +111,7 @@ private static void writePyproject(
                     reportUnusedFunction = false
                     reportUnusedVariable = false
                     reportUnnecessaryComparison = false
+                    reportUnusedClass = false
 
                     [tool.black]
                     target-version = ["py311"]

codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/StructureGenerator.java

@@ -20,13 +20,16 @@
 import java.time.ZonedDateTime;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.knowledge.HttpBinding;
 import software.amazon.smithy.model.knowledge.HttpBindingIndex;
 import software.amazon.smithy.model.knowledge.NullableIndex;
+import software.amazon.smithy.model.knowledge.OperationIndex;
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.Shape;
@@ -37,6 +40,7 @@
 import software.amazon.smithy.model.traits.InputTrait;
 import software.amazon.smithy.model.traits.OutputTrait;
 import software.amazon.smithy.model.traits.SensitiveTrait;
+import software.amazon.smithy.python.codegen.generators.MemberDeserializerGenerator;
 import software.amazon.smithy.python.codegen.generators.MemberSerializerGenerator;
 
 
@@ -109,10 +113,13 @@ class $L:
 
                     ${C|}
 
+                    ${C|}
+
                 """, symbol.getName(),
                 writer.consumer(w -> writeClassDocs(false)),
                 writer.consumer(w -> writeProperties()),
-                writer.consumer(w -> generateSerializeMethod()));
+                writer.consumer(w -> generateSerializeMethod()),
+                writer.consumer(w -> generateDeserializeMethod()));
     }
 
     private void renderError() {
@@ -334,20 +341,31 @@ def serialize(self, serializer: ShapeSerializer):
 
     private List<MemberShape> filterMembers() {
         var httpIndex = HttpBindingIndex.of(model);
+        var operationIndex = OperationIndex.of(model);
         if (shape.hasTrait(InputTrait.class)) {
-            var bindings = httpIndex.getRequestBindings(shape).keySet();
+            var operation = operationIndex.getInputBindings(shape).iterator().next();
+            var bindings = httpIndex.getRequestBindings(operation);
             return shape.members().stream()
-                    .filter(member -> bindings.contains(member.getMemberName()))
+                    .filter(member -> filterMember(member, bindings))
                     .toList();
         } else if (shape.hasTrait(OutputTrait.class)) {
-            var bindings = httpIndex.getResponseBindings(shape).keySet();
+            var operation = operationIndex.getOutputBindings(shape).iterator().next();
+            var bindings = httpIndex.getResponseBindings(operation);
             return shape.members().stream()
-                    .filter(member -> bindings.contains(member.getMemberName()))
+                    .filter(member -> filterMember(member, bindings))
                     .toList();
         }
         return shape.members().stream().toList();
     }
 
+    private boolean filterMember(MemberShape member, Map<String, HttpBinding> bindings) {
+        if (bindings.containsKey(member.getMemberName())) {
+            var binding = bindings.get(member.getMemberName());
+            return binding.getLocation() == HttpBinding.Location.DOCUMENT;
+        }
+        return true;
+    }
+
     private boolean isNullable(MemberShape member) {
         if (!NullableIndex.of(model).isMemberNullable(member)) {
             return false;
@@ -356,4 +374,45 @@ private boolean isNullable(MemberShape member) {
         return !target.isDocumentShape() || member.getMemberTrait(model, DefaultTrait.class).isEmpty();
     }
 
+    private void generateDeserializeMethod() {
+        writer.pushState();
+        writer.addLogger();
+        writer.addStdlibImports("typing", Set.of("Self", "Any"));
+        writer.addImport("smithy_core.deserializers", "ShapeDeserializer");
+
+        var deserializeableMembers = filterMembers();
+        var schemaSymbol = symbolProvider.toSymbol(shape).expectProperty(SymbolProperties.SCHEMA);
+        writer.putContext("schema", schemaSymbol);
+        writer.write("""
+                @classmethod
+                def deserialize(cls, deserializer: ShapeDeserializer) -> Self:
+                    kwargs: dict[str, Any] = {}
+
+                    def _consumer(schema: Schema, de: ShapeDeserializer) -> None:
+                        match schema.expect_member_index():
+                            ${C|}
+                            case _:
+                                logger.debug(f"Unexpected member schema: {schema}")
+
+                    deserializer.read_struct($T, consumer=_consumer)
+                    return cls(**kwargs)
+
+                """,
+                writer.consumer(w -> deserializeMembers(deserializeableMembers)),
+                schemaSymbol);
+        writer.popState();
+    }
+
+    private void deserializeMembers(List<MemberShape> members) {
+        int index = 0;
+        for (MemberShape member : members) {
+            var target = model.expectShape(member.getTarget());
+            writer.write("""
+                    case $L:
+                        kwargs[$S] = ${C|}
+                    """, index++, symbolProvider.toMemberName(member), writer.consumer(w ->
+                    target.accept(new MemberDeserializerGenerator(context, writer, member, "de"))
+            ));
+        }
+    }
 }

codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/SymbolProperties.java

@@ -68,5 +68,11 @@ public final class SymbolProperties {
      */
     public static final Property<Symbol> SERIALIZER = Property.named("serializer");
 
+    /**
+     * Contains a symbol pointing to the shape's deserializer method. This is only used for
+     * lists, maps, and unions.
+     */
+    public static final Property<Symbol> DESERIALIZER = Property.named("deserializer");
+
     private SymbolProperties() {}
 }

codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/SymbolVisitor.java

@@ -191,6 +191,10 @@ public Symbol listShape(ListShape shape) {
                 shape, "_serialize_" + CaseUtils.toSnakeCase(shape.getId().getName()), SHAPES_FILE, false)
                 .build());
 
+        builder.putProperty(SymbolProperties.DESERIALIZER, createGeneratedSymbolBuilder(
+                shape, "_deserialize_" + CaseUtils.toSnakeCase(shape.getId().getName()), SHAPES_FILE, false)
+                .build());
+
         return builder.build();
     }
 
@@ -206,6 +210,10 @@ public Symbol mapShape(MapShape shape) {
                 shape, "_serialize_" + CaseUtils.toSnakeCase(shape.getId().getName()), SHAPES_FILE, false)
                 .build());
 
+        builder.putProperty(SymbolProperties.DESERIALIZER, createGeneratedSymbolBuilder(
+                shape, "_deserialize_" + CaseUtils.toSnakeCase(shape.getId().getName()), SHAPES_FILE, false)
+                .build());
+
         return builder.build();
     }
 
@@ -332,9 +340,14 @@ public Symbol unionShape(UnionShape shape) {
 
         var unknownName = name + "Unknown";
         var unknownSymbol = createGeneratedSymbolBuilder(shape, unknownName, SHAPES_FILE).build();
-        return createGeneratedSymbolBuilder(shape, name, SHAPES_FILE)
-                .putProperty(SymbolProperties.UNION_UNKNOWN, unknownSymbol)
-                .build();
+        var builder = createGeneratedSymbolBuilder(shape, name, SHAPES_FILE)
+                .putProperty(SymbolProperties.UNION_UNKNOWN, unknownSymbol);
+
+        builder.putProperty(SymbolProperties.DESERIALIZER, createGeneratedSymbolBuilder(
+                shape, "_" + name + "Deserializer", SHAPES_FILE, false)
+                .build());
+
+        return builder.build();
     }
 
     @Override

codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/UnionGenerator.java

@@ -24,6 +24,7 @@
 import software.amazon.smithy.model.shapes.UnionShape;
 import software.amazon.smithy.model.traits.DocumentationTrait;
 import software.amazon.smithy.model.traits.StringTrait;
+import software.amazon.smithy.python.codegen.generators.MemberDeserializerGenerator;
 import software.amazon.smithy.python.codegen.generators.MemberSerializerGenerator;
 
 /**
@@ -54,6 +55,7 @@ final class UnionGenerator implements Runnable {
 
     @Override
     public void run() {
+        writer.pushState();
         var parentName = symbolProvider.toSymbol(shape).getName();
         writer.addStdlibImport("dataclasses", "dataclass");
         writer.addImport("smithy_core.serializers", "ShapeSerializer");
@@ -120,8 +122,64 @@ raise SmithyException("Unknown union variants may not be serialized.")
         memberNames.add(unknownSymbol.getName());
 
         shape.getTrait(DocumentationTrait.class).ifPresent(trait -> writer.writeComment(trait.getValue()));
-        writer.addStdlibImport("typing", "Union");
-        writer.write("$L = Union[$L]", parentName, String.join(", ", memberNames));
+        writer.write("type $L = $L\n", parentName, String.join(" | ", memberNames));
 
+        generateDeserializer();
+        writer.popState();
+    }
+
+    private void generateDeserializer() {
+        writer.addLogger();
+        writer.addStdlibImports("typing", Set.of("Self", "Any"));
+        writer.addImport("smithy_core.deserializers", "ShapeDeserializer");
+        writer.addImport("smithy_core.exceptions", "SmithyException");
+
+        // TODO: add in unknown handling
+
+        var symbol = symbolProvider.toSymbol(shape);
+        var deserializerSymbol = symbol.expectProperty(SymbolProperties.DESERIALIZER);
+        var schemaSymbol = symbol.expectProperty(SymbolProperties.SCHEMA);
+        writer.putContext("schema", schemaSymbol);
+        writer.write("""
+                class $1L:
+                    _result: $2T | None = None
+
+                    def deserialize(self, deserializer: ShapeDeserializer) -> $2T:
+                        self._result = None
+                        deserializer.read_struct($3T, self._consumer)
+
+                        if self._result is None:
+                            raise SmithyException("Unions must have exactly one value, but found none.")
+
+                        return self._result
+
+                    def _consumer(self, schema: Schema, de: ShapeDeserializer) -> None:
+                        match schema.expect_member_index():
+                            ${4C|}
+                            case _:
+                                logger.debug(f"Unexpected member schema: {schema}")
+
+                    def _set_result(self, value: $2T) -> None:
+                        if self._result is not None:
+                            raise SmithyException("Unions must have exactly one value, but found more than one.")
+                        self._result = value
+                """,
+                deserializerSymbol.getName(),
+                symbol,
+                schemaSymbol,
+                writer.consumer(w -> deserializeMembers()));
+    }
+
+    private void deserializeMembers() {
+        int index = 0;
+        for (MemberShape member : shape.members()) {
+            var target = model.expectShape(member.getTarget());
+            writer.write("""
+                    case $L:
+                        self._set_result($T(${C|}))
+                    """, index++, symbolProvider.toSymbol(member), writer.consumer(w ->
+                    target.accept(new MemberDeserializerGenerator(context, writer, member, "de"))
+            ));
+        }
     }
 }

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

@@ -30,6 +30,7 @@ public ListGenerator(GenerationContext context, PythonWriter writer, ListShape s
     @Override
     public void run() {
         generateSerializer();
+        generateDeserializer();
     }
 
     private void generateSerializer() {
@@ -61,4 +62,39 @@ private void generateSerializer() {
                         new MemberSerializerGenerator(context, w, shape.getMember(), "ls"))));
         writer.popState();
     }
+
+    private void generateDeserializer() {
+        var listSymbol = context.symbolProvider().toSymbol(shape);
+        var deserializerSymbol = listSymbol.expectProperty(SymbolProperties.DESERIALIZER);
+        var memberTarget = context.model().expectShape(shape.getMember().getTarget());
+
+        writer.pushState();
+        writer.addImport("smithy_core.serializers", "ShapeSerializer");
+        writer.addImport("smithy_core.schemas", "Schema");
+        var sparse = shape.hasTrait(SparseTrait.class);
+        writer.putContext("sparse", sparse);
+        writer.putContext("includeSchema", sparse || (
+                !memberTarget.isUnionShape() && !memberTarget.isStructureShape()));
+        writer.write("""
+                def $1L(deserializer: ShapeDeserializer, schema: Schema) -> $2T:
+                    result: $2T = []
+                    ${?includeSchema}
+                    member_schema = schema.members["member"]
+                    ${/includeSchema}
+                    deserializer.read_list(
+                        schema,
+                        ${?sparse}
+                        lambda d: result.append(d.read_optional(member_schema, lambda s: ${3C|}))
+                        ${/sparse}
+                        ${^sparse}
+                        lambda d: result.append(${3C|})
+                        ${/sparse}
+                    )
+                    return result
+                """, deserializerSymbol.getName(), listSymbol,
+                writer.consumer(w -> memberTarget.accept(
+                        new MemberDeserializerGenerator(context, w, shape.getMember(), "d")
+                )));
+        writer.popState();
+    }
 }

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

@@ -30,6 +30,7 @@ public MapGenerator(GenerationContext context, PythonWriter writer, MapShape sha
     @Override
     public void run() {
         generateSerializer();
+        generateDeserializer();
     }
 
     private void generateSerializer() {
@@ -64,4 +65,37 @@ private void generateSerializer() {
                         new MemberSerializerGenerator(context, w, shape.getValue(), "vs"))));
         writer.popState();
     }
+
+    private void generateDeserializer() {
+        var listSymbol = context.symbolProvider().toSymbol(shape);
+        var deserializerSymbol = listSymbol.expectProperty(SymbolProperties.DESERIALIZER);
+        var valueTarget = context.model().expectShape(shape.getValue().getTarget());
+
+        writer.pushState();
+        writer.addImport("smithy_core.serializers", "ShapeSerializer");
+        writer.addImport("smithy_core.schemas", "Schema");
+        var sparse = shape.hasTrait(SparseTrait.class);
+        writer.putContext("sparse", sparse);
+        writer.putContext("includeSchema", sparse || (
+                !valueTarget.isUnionShape() && !valueTarget.isStructureShape()));
+        writer.write("""
+                def $1L(deserializer: ShapeDeserializer, schema: Schema) -> $2T:
+                    result: $2T = {}
+                    value_schema = schema.members["value"]
+                    deserializer.read_map(
+                        schema,
+                        ${?sparse}
+                        lambda k, d: result.__setitem__(k, d.read_optional(value_schema, lambda s: ${3C|}))
+                        ${/sparse}
+                        ${^sparse}
+                        lambda k, d: result.__setitem__(k, ${3C|})
+                        ${/sparse}
+                    )
+                    return result
+                """, deserializerSymbol.getName(), listSymbol,
+                writer.consumer(w -> valueTarget.accept(
+                        new MemberDeserializerGenerator(context, w, shape.getValue(), "d")
+                )));
+        writer.popState();
+    }
 }