checkpoint 2

Author: haydenbaker

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

@@ -25,6 +25,7 @@
 import software.amazon.smithy.python.codegen.sections.SendRequestSection;
 import software.amazon.smithy.python.codegen.sections.SignRequestSection;
 import software.amazon.smithy.python.codegen.writer.PythonWriter;
+import software.amazon.smithy.utils.CaseUtils;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
 /**
@@ -705,6 +706,7 @@ private void writeDefaultPlugins(PythonWriter writer, Collection<SymbolReference
      */
     private void generateOperation(PythonWriter writer, OperationShape operation) {
         var operationSymbol = context.symbolProvider().toSymbol(operation);
+        var operationMethodSymbol = operationSymbol.expectProperty(SymbolProperties.OPERATION_METHOD);
         var pluginSymbol = CodegenUtils.getPluginSymbol(context.settings());
 
         var input = context.model().expectShape(operation.getInputShape());
@@ -715,7 +717,7 @@ private void generateOperation(PythonWriter writer, OperationShape operation) {
 
         writer.openBlock("async def $L(self, input: $T, plugins: list[$T] | None = None) -> $T:",
                 "",
-                operationSymbol.getName(),
+                operationMethodSymbol.getName(),
                 inputSymbol,
                 pluginSymbol,
                 outputSymbol,

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

@@ -21,6 +21,7 @@
 import software.amazon.smithy.codegen.core.directed.GenerateIntEnumDirective;
 import software.amazon.smithy.codegen.core.directed.GenerateListDirective;
 import software.amazon.smithy.codegen.core.directed.GenerateMapDirective;
+import software.amazon.smithy.codegen.core.directed.GenerateOperationDirective;
 import software.amazon.smithy.codegen.core.directed.GenerateServiceDirective;
 import software.amazon.smithy.codegen.core.directed.GenerateStructureDirective;
 import software.amazon.smithy.codegen.core.directed.GenerateUnionDirective;
@@ -35,6 +36,7 @@
 import software.amazon.smithy.python.codegen.generators.IntEnumGenerator;
 import software.amazon.smithy.python.codegen.generators.ListGenerator;
 import software.amazon.smithy.python.codegen.generators.MapGenerator;
+import software.amazon.smithy.python.codegen.generators.OperationGenerator;
 import software.amazon.smithy.python.codegen.generators.ProtocolGenerator;
 import software.amazon.smithy.python.codegen.generators.SchemaGenerator;
 import software.amazon.smithy.python.codegen.generators.ServiceErrorGenerator;
@@ -133,6 +135,19 @@ public void generateService(GenerateServiceDirective<GenerationContext, PythonSe
         protocolGenerator.generateProtocolTests(directive.context());
     }
 
+    @Override
+    public void generateOperation(GenerateOperationDirective<GenerationContext, PythonSettings> directive) {
+        DirectedCodegen.super.generateOperation(directive);
+
+        directive.context().writerDelegator().useShapeWriter(directive.shape(), writer -> {
+            OperationGenerator generator = new OperationGenerator(
+                    directive.context(),
+                    writer,
+                    directive.shape());
+            generator.run();
+        });
+    }
+
     @Override
     public void generateStructure(GenerateStructureDirective<GenerationContext, PythonSettings> directive) {
         directive.context().writerDelegator().useShapeWriter(directive.shape(), writer -> {

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

@@ -276,8 +276,14 @@ public Symbol bigDecimalShape(BigDecimalShape shape) {
     public Symbol operationShape(OperationShape shape) {
         // Operation names are escaped like members because ultimately they're
         // properties on an object too.
-        var name = escaper.escapeMemberName(CaseUtils.toSnakeCase(shape.getId().getName(service)));
-        return createGeneratedSymbolBuilder(shape, name, "client").build();
+        var methodName = escaper.escapeMemberName(CaseUtils.toSnakeCase(shape.getId().getName(service)));
+        var methodSymbol = createGeneratedSymbolBuilder(shape, methodName, "client").build();
+
+        // We add a symbol for the method in the client as a property, whereas the actual
+        // operation symbol points to the generated type for it
+        return createGeneratedSymbolBuilder(shape, getDefaultShapeName(shape), SHAPES_FILE)
+                .putProperty(SymbolProperties.OPERATION_METHOD, methodSymbol)
+                .build();
     }
 
     @Override

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

@@ -73,5 +73,11 @@ public final class SymbolProperties {
      */
     public static final Property<Symbol> DESERIALIZER = Property.named("deserializer");
 
+    /**
+     * Contains a symbol pointing to an operation shape's method in the client. This is
+     * only used for operations.
+     */
+    public static final Property<Symbol> OPERATION_METHOD = Property.named("operationMethod");
+
     private SymbolProperties() {}
 }

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

@@ -0,0 +1,104 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.codegen.generators;
+
+import software.amazon.smithy.codegen.core.Symbol;
+import software.amazon.smithy.codegen.core.SymbolProvider;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.knowledge.ServiceIndex;
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.python.codegen.CodegenUtils;
+import software.amazon.smithy.python.codegen.GenerationContext;
+import software.amazon.smithy.python.codegen.SymbolProperties;
+import software.amazon.smithy.python.codegen.writer.PythonWriter;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+import java.util.List;
+import java.util.logging.Logger;
+
+@SmithyInternalApi
+public final class OperationGenerator implements Runnable {
+    private static final Logger LOGGER = Logger.getLogger(OperationGenerator.class.getName());
+
+    private final GenerationContext context;
+    private final PythonWriter writer;
+    private final OperationShape shape;
+    private final SymbolProvider symbolProvider;
+    private final Model model;
+
+
+    public OperationGenerator(GenerationContext context, PythonWriter writer, OperationShape shape) {
+        this.context = context;
+        this.writer = writer;
+        this.shape = shape;
+        this.symbolProvider = context.symbolProvider();
+        this.model = context.model();
+    }
+
+    @Override
+    public void run() {
+        var opSymbol = symbolProvider.toSymbol(shape);
+        var inSymbol = symbolProvider.toSymbol(model.expectShape(shape.getInputShape()));
+        var outSymbol = symbolProvider.toSymbol(model.expectShape(shape.getOutputShape()));
+
+        writer.addStdlibImport("dataclasses", "dataclass");
+        writer.addImport("smithy_core.schemas", "APIOperation");
+        writer.addImport("smithy_core.type_registry", "TypeRegistry");
+
+        writer.write("""
+                @dataclass(kw_only=True, frozen=True)
+                class $1L(APIOperation["$2T", "$3T"]):
+                        input = $2T                        
+                        output = $3T
+                        schema = $4T
+                        input_schema = $5T
+                        output_schema = $6T
+                        error_registry = TypeRegistry({
+                            $7C
+                        })
+                        effective_auth_schemes = [
+                            $8C
+                        ]
+                """,
+                opSymbol.getName(),
+                inSymbol,
+                outSymbol,
+                opSymbol.expectProperty(SymbolProperties.SCHEMA),
+                inSymbol.expectProperty(SymbolProperties.SCHEMA),
+                outSymbol.expectProperty(SymbolProperties.SCHEMA),
+                writer.consumer(this::writeErrorTypeRegistry),
+                writer.consumer(this::writeAuthSchemes)
+                // TODO: Docs? Maybe not necessary on the operation type itself
+                // TODO: Singleton?
+        );
+    }
+
+    private void writeErrorTypeRegistry(PythonWriter writer) {
+        List<ShapeId> errors = shape.getErrors();
+        if (!errors.isEmpty()) {
+            writer.addImport("smithy_core.shapes", "ShapeID");
+        }
+        for (var error : errors) {
+            var errSymbol = symbolProvider.toSymbol(model.expectShape(error));
+            writer.write("ShapeID($S): $T,", error, errSymbol);
+        }
+    }
+
+    private void writeAuthSchemes(PythonWriter writer) {
+        var authSchemes = ServiceIndex.of(model).getEffectiveAuthSchemes(context.settings().service(), shape.getId(),
+                ServiceIndex.AuthSchemeMode.NO_AUTH_AWARE);
+
+        if (!authSchemes.isEmpty()) {
+            writer.addImport("smithy_core.shapes", "ShapeID");
+        }
+
+        for(var authSchemeId : authSchemes.keySet()) {
+            writer.write("ShapeID($S)", authSchemeId);
+        }
+
+    }
+}

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

@@ -146,8 +146,7 @@ def shape_type(self) -> ShapeType:
     @property
     def discriminator(self) -> ShapeID:
         """The shape ID that corresponds to the contents of the document."""
-        # TODO: custom exception?
-        raise NotImplementedError(f"{self} document has no discriminator.")
+        return self._schema.id
 
     def is_none(self) -> bool:
         """Indicates whether the document contains a null value."""

packages/smithy-core/src/smithy_core/schemas.py

@@ -7,7 +7,7 @@
 from .exceptions import ExpectationNotMetException, SmithyException
 from .shapes import ShapeID, ShapeType
 from .traits import Trait, DynamicTrait, IdempotencyTokenTrait, StreamingTrait
-
+from .type_registry import TypeRegistry
 
 if TYPE_CHECKING:
     from .serializers import SerializeableShape
@@ -289,8 +289,7 @@ class APIOperation[I: "SerializeableShape", O: "DeserializeableShape"]:
     output_schema: Schema
     """The schema of the operation's output shape."""
 
-    # TODO: Add a type registry for errors
-    error_registry: Any
+    error_registry: TypeRegistry
     """A TypeRegistry used to create errors."""
 
     effective_auth_schemes: Sequence[ShapeID]

packages/smithy-core/src/smithy_core/type_registry.py

@@ -13,18 +13,18 @@
 class TypeRegistry:
     def __init__(
         self,
-        types: dict[ShapeID, DeserializeableShape],
+        types: dict[ShapeID, type[DeserializeableShape]],
         sub_registry: "TypeRegistry | None" = None,
     ):
         self._types = types
         self._sub_registry = sub_registry
 
     def get(self, shape: ShapeID) -> type[DeserializeableShape]:
         if shape in self._types:
-            return type(self._types[shape])
+            return self._types[shape]
         if self._sub_registry is not None:
             return self._sub_registry.get(shape)
-        raise ValueError(f"Unknown shape: {shape}")  # TODO: real exception?
+        raise KeyError(f"Unknown shape: {shape}")
 
     def deserialize(self, document: Document) -> DeserializeableShape:
         return document.as_shape(self.get(document.discriminator))

packages/smithy-core/tests/unit/test_type_registry.py

@@ -0,0 +1,48 @@
+from smithy_core.deserializers import DeserializeableShape, ShapeDeserializer
+from smithy_core.documents import Document
+from smithy_core.schemas import Schema
+from smithy_core.shapes import ShapeID, ShapeType
+from smithy_core.type_registry import TypeRegistry
+import pytest
+
+
+class TestTypeRegistry:
+    def test_get(self):
+        registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+
+        result = registry.get(ShapeID("com.example#Test"))
+
+        assert result == TestShape
+
+    def test_get_sub_registry(self):
+        sub_registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+        registry = TypeRegistry({}, sub_registry)
+
+        result = registry.get(ShapeID("com.example#Test"))
+
+        assert result == TestShape
+
+    def test_get_no_match(self):
+        registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+
+        with pytest.raises(KeyError, match="Unknown shape: com.example#Test2"):
+            registry.get(ShapeID("com.example#Test2"))
+
+    def test_deserialize(self):
+        shape_id = ShapeID("com.example#Test")
+        registry = TypeRegistry({shape_id: TestShape})
+
+        result = registry.deserialize(Document("abc123", schema=TestShape.schema))
+
+        assert isinstance(result, TestShape) and result.value == "abc123"
+
+
+class TestShape(DeserializeableShape):
+    schema = Schema(id=ShapeID("com.example#Test"), shape_type=ShapeType.STRING)
+
+    def __init__(self, value: str):
+        self.value = value
+
+    @classmethod
+    def deserialize(cls, deserializer: ShapeDeserializer) -> "TestShape":
+        return TestShape(deserializer.read_string(schema=TestShape.schema))