Implement client error correction

codegen/build.gradle.kts

@@ -15,5 +15,5 @@
 
 allprojects {
     group = "software.amazon.smithy.python"
-    version = "0.3.0"
+    version = "0.3.1"
 }

codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java

@@ -31,6 +31,7 @@
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.traits.DefaultTrait;
 import software.amazon.smithy.model.traits.ErrorTrait;
 import software.amazon.smithy.model.traits.TimestampFormatTrait;
 import software.amazon.smithy.model.traits.TimestampFormatTrait.Format;
@@ -128,6 +129,18 @@ public static boolean isErrorMessage(Model model, MemberShape shape) {
                 && model.expectShape(shape.getContainer()).hasTrait(ErrorTrait.class);
     }
 
+    /**
+     * Determines whether a member is required in the generated Python constructor — that is,
+     * neither nullable nor carrying a default value.
+     *
+     * @param index A nullable index for the model.
+     * @param member The member to check.
+     * @return Returns whether the member is required in generated code.
+     */
+    public static boolean isRequiredMember(NullableIndex index, MemberShape member) {
+        return !index.isMemberNullable(member) && !member.hasTrait(DefaultTrait.class);
+    }
+
     /**
      * Executes a given shell command in a given directory.
      *

codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java

@@ -763,12 +763,15 @@ public Void nullNode(NullNode node) {
 
         @Override
         public Void numberNode(NumberNode node) {
-            // TODO: Add support for timestamp, int-enum, and others
             if (inputShape.isTimestampShape()) {
                 var parsed = CodegenUtils.parseTimestampNode(model, inputShape, node);
                 writer.writeInline(CodegenUtils.getDatetimeConstructor(writer, parsed));
             } else if (inputShape.isFloatShape() || inputShape.isDoubleShape()) {
                 writer.writeInline("float($L)", node.getValue());
+            } else if (inputShape.isIntEnumShape()) {
+                var enumSymbol =
+                        context.symbolProvider().toSymbol(inputShape);
+                writer.writeInline("$T($L)", enumSymbol, node.getValue());
             } else {
                 writer.writeInline("$L", node.getValue());
             }
@@ -800,6 +803,10 @@ public Void stringNode(StringNode node) {
                 };
 
                 writer.writeInline("float($S)", value);
+            } else if (inputShape.isEnumShape()) {
+                var enumSymbol =
+                        context.symbolProvider().toSymbol(inputShape);
+                writer.writeInline("$T($S)", enumSymbol, node.getValue());
             } else {
                 writer.writeInline("$S", node.getValue());
             }

codegen/core/src/main/java/software/amazon/smithy/python/codegen/PythonSymbolProvider.java

@@ -329,16 +329,8 @@ public Symbol intEnumShape(IntEnumShape shape) {
     }
 
     private Symbol genericEnum(Shape shape) {
-        var enumSymbol = createGeneratedSymbolBuilder(shape, getDefaultShapeName(shape), SHAPES_FILE).build();
-
-        // We add this enum symbol as a property on a generic string/int symbol
-        // rather than returning the enum symbol directly because we only
-        // generate the enum constants for convenience. We actually want
-        // to pass around plain types rather than what is effectively
-        // a namespace class.
-        return createSymbolBuilder(shape, shape.isEnumShape() ? "str" : "int")
-                .putProperty(SymbolProperties.ENUM_SYMBOL, escaper.escapeSymbol(shape, enumSymbol))
-                .build();
+        Symbol symbol = createGeneratedSymbolBuilder(shape, getDefaultShapeName(shape), SHAPES_FILE).build();
+        return escaper.escapeSymbol(shape, symbol);
     }
 
     @Override

codegen/core/src/main/java/software/amazon/smithy/python/codegen/RequiredMemberTargetIndex.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.codegen;
+
+import java.util.HashSet;
+import java.util.Set;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.knowledge.KnowledgeIndex;
+import software.amazon.smithy.model.knowledge.NullableIndex;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+/**
+ * Knowledge index of all shapes that are the target of a python-required structure member,
+ * i.e. the shapes that may need a synthesized default during client error correction.
+ *
+ * <p>Computed once per model and cached via {@link Model#getKnowledge}.
+ */
+@SmithyInternalApi
+public final class RequiredMemberTargetIndex implements KnowledgeIndex {
+
+    private final Set<ShapeId> targets = new HashSet<>();
+
+    private RequiredMemberTargetIndex(Model model) {
+        var index = NullableIndex.of(model);
+        for (var struct : model.getStructureShapes()) {
+            for (var member : struct.members()) {
+                if (CodegenUtils.isRequiredMember(index, member)) {
+                    targets.add(member.getTarget());
+                }
+            }
+        }
+    }
+
+    public static RequiredMemberTargetIndex of(Model model) {
+        return model.getKnowledge(RequiredMemberTargetIndex.class, RequiredMemberTargetIndex::new);
+    }
+
+    /**
+     * @param shape The shape to check.
+     * @return Returns whether the shape is the target of any python-required structure member.
+     */
+    public boolean isRequiredMemberTarget(ShapeId shape) {
+        return targets.contains(shape);
+    }
+}

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

@@ -35,13 +35,6 @@ public final class SymbolProperties {
      */
     public static final Property<Boolean> STDLIB = Property.named("stdlib");
 
-    /**
-     * Contains a symbol representing the class containing known enum values for the symbol.
-     *
-     * <p>In type signatures, the base int or str is used instead for forwards compatibility.
-     */
-    public static final Property<Symbol> ENUM_SYMBOL = Property.named("enumSymbol");
-
     /**
      * Contains a symbol representing the unknown variant of a union.
      */

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

@@ -9,11 +9,19 @@
 import software.amazon.smithy.model.traits.DocumentationTrait;
 import software.amazon.smithy.model.traits.EnumValueTrait;
 import software.amazon.smithy.python.codegen.GenerationContext;
-import software.amazon.smithy.python.codegen.SymbolProperties;
+import software.amazon.smithy.python.codegen.SmithyPythonDependency;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
 /**
- * Renders enums.
+ * Renders enums as a {@code StrEnum} subclass.
+ *
+ * <p>The {@code UnknownEnumMixin} base from smithy-core handles values that
+ * weren't known at generation time: deserializing an unrecognized value (or
+ * filling a missing required member during client error correction, see
+ * {@link MemberErrorCorrectionGenerator}) produces a pseudo-member with
+ * {@code is_unknown} set rather than raising.
+ *
+ * @see <a href="https://smithy.io/2.0/spec/simple-types.html#enum">Smithy spec: enum</a>
  */
 @SmithyInternalApi
 public final class EnumGenerator implements Runnable {
@@ -27,10 +35,12 @@ public EnumGenerator(GenerationContext context, EnumShape enumShape) {
 
     @Override
     public void run() {
-        var enumSymbol = context.symbolProvider().toSymbol(shape).expectProperty(SymbolProperties.ENUM_SYMBOL);
+        var enumSymbol = context.symbolProvider().toSymbol(shape);
         context.writerDelegator().useShapeWriter(shape, writer -> {
             writer.addStdlibImport("enum", "StrEnum");
-            writer.openBlock("class $L(StrEnum):", "", enumSymbol.getName(), () -> {
+            writer.addDependency(SmithyPythonDependency.SMITHY_CORE);
+            writer.addImport("smithy_core.types", "UnknownEnumMixin");
+            writer.openBlock("class $L(UnknownEnumMixin, StrEnum):", "", enumSymbol.getName(), () -> {
                 shape.getTrait(DocumentationTrait.class).ifPresent(trait -> {
                     writer.writeDocs(trait.getValue(), context);
                 });

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

@@ -10,9 +10,20 @@
 import software.amazon.smithy.model.traits.EnumValueTrait;
 import software.amazon.smithy.python.codegen.GenerationContext;
 import software.amazon.smithy.python.codegen.PythonSettings;
-import software.amazon.smithy.python.codegen.SymbolProperties;
+import software.amazon.smithy.python.codegen.SmithyPythonDependency;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
+/**
+ * Renders intEnums as an {@code IntEnum} subclass.
+ *
+ * <p>The {@code UnknownEnumMixin} base from smithy-core handles values that
+ * weren't known at generation time: deserializing an unrecognized value (or
+ * filling a missing required member during client error correction, see
+ * {@link MemberErrorCorrectionGenerator}) produces a pseudo-member with
+ * {@code is_unknown} set rather than raising.
+ *
+ * @see <a href="https://smithy.io/2.0/spec/simple-types.html#intenum">Smithy spec: intEnum</a>
+ */
 @SmithyInternalApi
 public final class IntEnumGenerator implements Runnable {
 
@@ -24,18 +35,20 @@ public IntEnumGenerator(GenerateIntEnumDirective<GenerationContext, PythonSettin
 
     @Override
     public void run() {
-        var enumSymbol = directive.symbol().expectProperty(SymbolProperties.ENUM_SYMBOL);
+        var enumSymbol = directive.symbol();
         directive.context().writerDelegator().useShapeWriter(directive.shape(), writer -> {
             writer.addStdlibImport("enum", "IntEnum");
-            writer.openBlock("class $L(IntEnum):", "", enumSymbol.getName(), () -> {
+            writer.addDependency(SmithyPythonDependency.SMITHY_CORE);
+            writer.addImport("smithy_core.types", "UnknownEnumMixin");
+            writer.openBlock("class $L(UnknownEnumMixin, IntEnum):", "", enumSymbol.getName(), () -> {
                 directive.shape().getTrait(DocumentationTrait.class).ifPresent(trait -> {
                     writer.writeDocs(trait.getValue(), directive.context());
                 });
 
                 for (MemberShape member : directive.shape().members()) {
                     var name = directive.symbolProvider().toMemberName(member);
                     var value = member.expectTrait(EnumValueTrait.class).expectIntValue();
-                    writer.write("$L = $L\n", name, value);
+                    writer.write("$L = $L", name, value);
                     member.getTrait(DocumentationTrait.class).ifPresent(trait -> {
                         writer.writeDocs(trait.getValue(), directive.context());
                     });

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

@@ -135,7 +135,11 @@ public Void integerShape(IntegerShape shape) {
 
     @Override
     public Void intEnumShape(IntEnumShape shape) {
-        writeDeserializer("integer");
+        pushMemberState();
+        var enumSymbol = context.symbolProvider().toSymbol(shape);
+        writer.write("$T(${deserializer:L}.read_integer(${C|}))",
+                enumSymbol,
+                writer.consumer(w -> writeSchema()));
         return null;
     }
 
@@ -183,7 +187,11 @@ public Void stringShape(StringShape shape) {
 
     @Override
     public Void enumShape(EnumShape shape) {
-        writeDeserializer("string");
+        pushMemberState();
+        var enumSymbol = context.symbolProvider().toSymbol(shape);
+        writer.write("$T(${deserializer:L}.read_string(${C|}))",
+                enumSymbol,
+                writer.consumer(w -> writeSchema()));
         return null;
     }