Generate event stream operation method signatures

This updates operation generation to generate event stream operations
with the EventStream type as its return value.

It also updates union generation so that unions contain their own
deserialize functions. This is needed to make the them pass the type
check, but also it is best to have them own as much of that as
possible so that the deserializer function can be left to only
dispatch duty.

Author: JordonPhillips

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

@@ -20,6 +20,8 @@
 import java.util.LinkedHashSet;
 import java.util.Set;
 import software.amazon.smithy.codegen.core.SymbolReference;
+import software.amazon.smithy.model.knowledge.EventStreamIndex;
+import software.amazon.smithy.model.knowledge.EventStreamInfo;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
 import software.amazon.smithy.model.knowledge.TopDownIndex;
 import software.amazon.smithy.model.shapes.OperationShape;
@@ -104,8 +106,14 @@ def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None):
                     """, configSymbol, pluginSymbol, writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)));
 
             var topDownIndex = TopDownIndex.of(context.model());
+            var eventStreamIndex = EventStreamIndex.of(context.model());
             for (OperationShape operation : topDownIndex.getContainedOperations(service)) {
-                generateOperation(writer, operation);
+                if (eventStreamIndex.getInputInfo(operation).isPresent()
+                        || eventStreamIndex.getOutputInfo(operation).isPresent()) {
+                    generateEventStreamOperation(writer, operation);
+                } else {
+                    generateOperation(writer, operation);
+                }
             }
         });
 
@@ -348,7 +356,7 @@ async def _handle_attempt(
                         )
 
                 """, CodegenUtils.getHttpAuthParamsSymbol(context.settings()),
-                writer.consumer(this::initializeHttpAuthParameters));
+                    writer.consumer(this::initializeHttpAuthParameters));
             writer.popState();
 
             writer.addDependency(SmithyPythonDependency.SMITHY_CORE);
@@ -641,48 +649,48 @@ private void generateOperation(PythonWriter writer, OperationShape operation) {
 
         writer.openBlock("async def $L(self, input: $T, plugins: list[$T] | None = None) -> $T:", "",
                 operationSymbol.getName(), inputSymbol, pluginSymbol, outputSymbol, () -> {
-            writer.writeDocs(() -> {
-                var docs = operation.getTrait(DocumentationTrait.class)
-                        .map(StringTrait::getValue)
-                        .orElse(String.format("Invokes the %s operation.", operation.getId().getName()));
+                    writer.writeDocs(() -> {
+                        var docs = operation.getTrait(DocumentationTrait.class)
+                                .map(StringTrait::getValue)
+                                .orElse(String.format("Invokes the %s operation.", operation.getId().getName()));
 
-                var inputDocs = input.getTrait(DocumentationTrait.class)
-                        .map(StringTrait::getValue)
-                        .orElse("The operation's input.");
+                        var inputDocs = input.getTrait(DocumentationTrait.class)
+                                .map(StringTrait::getValue)
+                                .orElse("The operation's input.");
 
-                writer.write("""
+                        writer.write("""
                         $L
 
                         :param input: $L
 
                         :param plugins: A list of callables that modify the configuration dynamically.
                         Changes made by these plugins only apply for the duration of the operation
                         execution and will not affect any other operation invocations.""", docs, inputDocs);
-            });
-
-            var defaultPlugins = new LinkedHashSet<SymbolReference>();
-            for (PythonIntegration integration : context.integrations()) {
-                for (RuntimeClientPlugin runtimeClientPlugin : integration.getClientPlugins()) {
-                    if (runtimeClientPlugin.matchesOperation(context.model(), service, operation)) {
-                        runtimeClientPlugin.getPythonPlugin().ifPresent(defaultPlugins::add);
+                    });
+
+                    var defaultPlugins = new LinkedHashSet<SymbolReference>();
+                    for (PythonIntegration integration : context.integrations()) {
+                        for (RuntimeClientPlugin runtimeClientPlugin : integration.getClientPlugins()) {
+                            if (runtimeClientPlugin.matchesOperation(context.model(), service, operation)) {
+                                runtimeClientPlugin.getPythonPlugin().ifPresent(defaultPlugins::add);
+                            }
+                        }
                     }
-                }
-            }
-            writer.write("""
+                    writer.write("""
                 operation_plugins: list[Plugin] = [
                     $C
                 ]
                 if plugins:
                     operation_plugins.extend(plugins)
                 """, writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)));
 
-            if (context.protocolGenerator() == null) {
-                writer.write("raise NotImplementedError()");
-            } else {
-                var protocolGenerator = context.protocolGenerator();
-                var serSymbol = protocolGenerator.getSerializationFunction(context, operation);
-                var deserSymbol = protocolGenerator.getDeserializationFunction(context, operation);
-                writer.write("""
+                    if (context.protocolGenerator() == null) {
+                        writer.write("raise NotImplementedError()");
+                    } else {
+                        var protocolGenerator = context.protocolGenerator();
+                        var serSymbol = protocolGenerator.getSerializationFunction(context, operation);
+                        var deserSymbol = protocolGenerator.getDeserializationFunction(context, operation);
+                        writer.write("""
                     return await self._execute_operation(
                         input=input,
                         plugins=operation_plugins,
@@ -692,7 +700,47 @@ private void generateOperation(PythonWriter writer, OperationShape operation) {
                         operation_name=$S,
                     )
                     """, serSymbol, deserSymbol, operation.getId().getName());
-            }
-        });
+                    }
+                });
+    }
+
+    private void generateEventStreamOperation(PythonWriter writer, OperationShape operation) {
+        writer.pushState();
+        writer.addDependency(SmithyPythonDependency.SMITHY_EVENT_STREAM);
+        writer.addImports("smithy_event_stream.aio.interfaces", Set.of(
+                "EventStream", "InputEventStream", "OutputEventStream"));
+        var operationSymbol = context.symbolProvider().toSymbol(operation);
+        var pluginSymbol = CodegenUtils.getPluginSymbol(context.settings());
+
+        var input = context.model().expectShape(operation.getInputShape());
+        var inputSymbol = context.symbolProvider().toSymbol(input);
+
+        var eventStreamIndex = EventStreamIndex.of(context.model());
+        var inputStreamSymbol = eventStreamIndex.getInputInfo(operation)
+                .map(EventStreamInfo::getEventStreamTarget)
+                .map(target -> context.symbolProvider().toSymbol(target))
+                .orElse(null);
+        writer.putContext("inputStream", inputStreamSymbol);
+
+        var output = context.model().expectShape(operation.getOutputShape());
+        var outputSymbol = context.symbolProvider().toSymbol(output);
+        var outputStreamSymbol = eventStreamIndex.getOutputInfo(operation)
+                .map(EventStreamInfo::getEventStreamTarget)
+                .map(target -> context.symbolProvider().toSymbol(target))
+                .orElse(null);
+        writer.putContext("outputStream", outputStreamSymbol);
+
+        writer.write("""
+                async def $L(self, input: $T, plugins: list[$T] | None = None) -> EventStream[
+                    ${?inputStream}InputEventStream[${inputStream:T}]${/inputStream}\
+                    ${^inputStream}None${/inputStream},
+                    ${?outputStream}OutputEventStream[${outputStream:T}]${/outputStream}\
+                    ${^outputStream}None${/outputStream},
+                    $T
+                ]:
+                    raise NotImplementedError()
+                """, operationSymbol.getName(), inputSymbol, pluginSymbol, outputSymbol);
+
+        writer.popState();
     }
 }

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

@@ -64,6 +64,13 @@ public final class SmithyPythonDependency {
             false
     );
 
+    public static final PythonDependency SMITHY_EVENT_STREAM = new PythonDependency(
+            "smithy_event_stream",
+            "==0.0.1",
+            Type.DEPENDENCY,
+            false
+    );
+
     /**
      * testing framework used in generated functional tests.
      */

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

@@ -72,16 +72,20 @@ public void run() {
 
             writer.write("""
                     @dataclass
-                    class $L:
-                        ${C|}
+                    class $1L:
+                        ${2C|}
 
-                        value: $T
+                        value: $3T
 
                         def serialize(self, serializer: ShapeSerializer):
-                            serializer.write_struct($T, self)
+                            serializer.write_struct($4T, self)
 
                         def serialize_members(self, serializer: ShapeSerializer):
-                            ${C|}
+                            ${5C|}
+
+                        @classmethod
+                        def deserialize(cls, deserializer: ShapeDeserializer) -> Self:
+                            return cls(value=${6C|})
 
                     """,
                     memberSymbol.getName(),
@@ -90,7 +94,11 @@ def serialize_members(self, serializer: ShapeSerializer):
                     targetSymbol,
                     schemaSymbol,
                     writer.consumer(w -> target.accept(
-                            new MemberSerializerGenerator(context, w, member, "serializer"))));
+                            new MemberSerializerGenerator(context, w, member, "serializer"))),
+                    writer.consumer(w -> target.accept(
+                            new MemberDeserializerGenerator(context, w, member, "deserializer")))
+
+            );
         }
 
         // Note that the unknown variant doesn't implement __eq__. This is because
@@ -118,11 +126,15 @@ raise SmithyException("Unknown union variants may not be serialized.")
                     def serialize_members(self, serializer: ShapeSerializer):
                         raise SmithyException("Unknown union variants may not be serialized.")
 
+                    @classmethod
+                    def deserialize(cls, deserializer: ShapeDeserializer) -> Self:
+                        raise NotImplementedError()
+
                 """, unknownSymbol.getName());
         memberNames.add(unknownSymbol.getName());
 
-        shape.getTrait(DocumentationTrait.class).ifPresent(trait -> writer.writeComment(trait.getValue()));
         writer.write("type $L = $L\n", parentName, String.join(" | ", memberNames));
+        shape.getTrait(DocumentationTrait.class).ifPresent(trait -> writer.writeDocs(trait.getValue()));
 
         generateDeserializer();
         writer.popState();
@@ -173,13 +185,10 @@ raise SmithyException("Unions must have exactly one value, but found more than o
     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"))
-            ));
+                        self._set_result($T.deserialize(de))
+                    """, index++, symbolProvider.toSymbol(member));
         }
     }
 }