Add protocol implementations, operation shape generation, type registry

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;
@@ -108,9 +110,10 @@ public void customizeBeforeShapeGeneration(CustomizeDirective<GenerationContext,
 
         new ConfigGenerator(directive.settings(), directive.context()).run();
         var serviceIndex = ServiceIndex.of(directive.model());
-        if (directive.context().applicationProtocol().isHttpProtocol()
-                && !serviceIndex.getAuthSchemes(directive.service()).isEmpty()) {
-            new HttpAuthGenerator(directive.context(), directive.settings()).run();
+        if (directive.context().applicationProtocol().isHttpProtocol()) {
+            if (!serviceIndex.getAuthSchemes(directive.service()).isEmpty()) {
+                new HttpAuthGenerator(directive.context(), directive.settings()).run();
+            }
             new EndpointsGenerator(directive.context(), directive.settings()).run();
         }
     }
@@ -133,6 +136,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", false).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,101 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.codegen.generators;
+
+import java.util.List;
+import java.util.logging.Logger;
+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.ShapeId;
+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;
+
+@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);
+        }
+
+    }
+}

codegen/core/src/main/java/software/amazon/smithy/python/codegen/writer/PythonWriter.java

@@ -21,10 +21,8 @@
 import software.amazon.smithy.model.node.NumberNode;
 import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.node.StringNode;
-import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.python.codegen.CodegenUtils;
 import software.amazon.smithy.python.codegen.PythonSettings;
-import software.amazon.smithy.python.codegen.SymbolProperties;
 import software.amazon.smithy.utils.SmithyUnstableApi;
 import software.amazon.smithy.utils.StringUtils;
 
@@ -305,9 +303,7 @@ public String apply(Object type, String indent) {
                 Symbol typeSymbol = (Symbol) type;
                 // Check if the symbol is an operation - we shouldn't add imports for operations, since
                 //  they are methods of the service object and *can't* be imported
-                if (!isOperationSymbol(typeSymbol)) {
-                    addUseImports(typeSymbol);
-                }
+                addUseImports(typeSymbol);
                 return typeSymbol.getName();
             } else if (type instanceof SymbolReference) {
                 SymbolReference typeSymbol = (SymbolReference) type;
@@ -320,10 +316,6 @@ public String apply(Object type, String indent) {
         }
     }
 
-    private Boolean isOperationSymbol(Symbol typeSymbol) {
-        return typeSymbol.getProperty(SymbolProperties.SHAPE).map(Shape::isOperationShape).orElse(false);
-    }
-
     private final class PythonNodeFormatter implements BiFunction<Object, String, String> {
         private final PythonWriter writer;
 

examples/weather/model/weather.smithy

@@ -26,6 +26,21 @@ resource City {
 resource Forecast {
     identifiers: { cityId: CityId }
     read: GetForecast,
+    update: PutForecast
+}
+
+@http(method: "PUT", uri: "/city/{cityId}/forecast", code: 200)
+@idempotent
+operation PutForecast {
+    input := for Forecast {
+        @required
+        @httpLabel
+        $cityId
+
+        chanceOfRain: Float
+    }
+
+    output := {}
 }
 
 // "pattern" is a trait.
@@ -154,3 +169,5 @@ structure GetForecastInput {
 structure GetForecastOutput {
     chanceOfRain: Float
 }
+
+

packages/smithy-core/src/smithy_core/aio/client.py

@@ -468,7 +468,7 @@ async def _handle_attempt[I: SerializeableShape, O: DeserializeableShape](
                 operation=call.operation,
                 request=response_context.transport_request,
                 response=response_context.transport_response,
-                error_registry="foo",
+                error_registry=call.operation.error_registry,
                 context=response_context.properties,
             )
 

packages/smithy-core/src/smithy_core/aio/interfaces/__init__.py

@@ -1,14 +1,14 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
 from collections.abc import AsyncIterable
-from typing import Protocol, runtime_checkable, TYPE_CHECKING, Any, Callable
+from typing import Protocol, runtime_checkable, TYPE_CHECKING, Callable
 
 from ...exceptions import UnsupportedStreamException
 from ...interfaces import URI, Endpoint, TypedProperties
 from ...interfaces import StreamingBlob as SyncStreamingBlob
 
 from .eventstream import EventPublisher, EventReceiver
-
+from ...type_registry import TypeRegistry
 
 if TYPE_CHECKING:
     from ...schemas import APIOperation
@@ -129,7 +129,7 @@ async def deserialize_response[
         operation: "APIOperation[OperationInput, OperationOutput]",
         request: I,
         response: O,
-        error_registry: Any,  # TODO: add error registry
+        error_registry: TypeRegistry,
         context: TypedProperties,
     ) -> OperationOutput:
         """Deserializes the output from the tranport response or throws an exception.

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

@@ -143,6 +143,11 @@ def shape_type(self) -> ShapeType:
         """The Smithy data model type for the underlying contents of the document."""
         return self._type
 
+    @property
+    def discriminator(self) -> ShapeID:
+        """The shape ID that corresponds to the contents of the document."""
+        return self._schema.id
+
     def is_none(self) -> bool:
         """Indicates whether the document contains a null value."""
         return self._value is None and self._raw_value is None