python: Update codegen templates to support struct enums (#1766)

Author: svix-jplatte

python/templates/component_type.py.jinja

@@ -1,8 +1,10 @@
 # this file is @generated
 {% if type.kind == "struct" -%}
     {% include "types/struct.py.jinja" -%}
-{%- elif type.kind == "string_enum"-%}
+{% elif type.kind == "string_enum" -%}
     {% include "types/string_enum.py.jinja" -%}
-{%- elif type.kind == "integer_enum" %}
-{%- include "types/integer_enum.py.jinja" -%}
-{%- endif %}
+{% elif type.kind == "integer_enum" -%}
+    {% include "types/integer_enum.py.jinja" -%}
+{% elif type.kind == "struct_enum" -%}
+    {% include "types/struct_enum.py.jinja" -%}
+{% endif %}

python/templates/types/struct.py.jinja

@@ -6,33 +6,10 @@ from .common import BaseModel
 
 {% for c in referenced_components -%}
 from . {{ c | to_snake_case }} import {{ c | to_upper_camel_case }}
-{% endfor -%}
+{% endfor %}
 
 class {{ type.name | to_upper_camel_case }}(BaseModel):
 {%- if type.description is defined %}
     {{ type.description | to_doc_comment(style="python") | indent(4) }}
 {% endif %}
-{%- for field in type.fields %}
-    {%- if field.required and not field.nullable %}
-        {%- if field.name | to_lower_camel_case != field.name %}
-    {{ field.name | to_snake_case }}: {{ field.type.to_python() }} = Field(alias="{{ field.name }}")
-        {%- else %}
-    {{ field.name | to_snake_case }}: {{ field.type.to_python() }}
-        {%- endif %}
-    {%- elif field.required and field.nullable %}
-        {%- if field.name | to_lower_camel_case != field.name %}
-    {{ field.name | to_snake_case }}: t.Optional[{{ field.type.to_python() }}] = Field(alias="{{ field.name }}")
-        {%- else %}
-    {{ field.name | to_snake_case }}: t.Optional[{{ field.type.to_python() }}]
-        {%- endif %}
-    {%- else %}
-        {%- if field.name | to_lower_camel_case != field.name %}
-    {{ field.name | to_snake_case }}: t.Optional[{{ field.type.to_python() }}] = Field(default=None, alias="{{ field.name }}")
-        {%- else %}
-    {{ field.name | to_snake_case }}: t.Optional[{{ field.type.to_python() }}] = None
-        {%- endif %}
-    {%- endif %}
-    {%- if field.description is defined %}
-    {{ field.description | to_doc_comment(style="python") | indent(4) }}
-    {%- endif %}
-{% endfor %}
+{%- include "types/struct_fields.py.jinja" %}

python/templates/types/struct_enum.py.jinja

@@ -0,0 +1,39 @@
+import typing as t
+from typing_extensions import Annotated
+
+from pydantic import Field
+from datetime import datetime
+
+from .common import BaseModel
+
+{% for c in referenced_components %}
+from . {{ c | to_snake_case }} import {{ c | to_upper_camel_case }}
+{%- endfor %}
+
+{%- set type_name = type.name | to_upper_camel_case %}
+
+
+{% for variant in type.variants %}
+class {{ type_name }}{{ variant.name | to_upper_camel_case }}(BaseModel):
+    {{ type.discriminator_field | to_snake_case }}: t.Literal["{{ variant.name }}"]
+
+    {% include "types/struct_fields.py.jinja" %}
+    {% if variant.schema_ref is defined -%}
+    {{ type.content_field | to_snake_case }}: {{ variant.schema_ref | to_upper_camel_case -}}
+        {% if type.content_field | to_lower_camel_case != type.content_field %}= Field(alias="{{ type.content_field }}"){% endif %}
+    {% endif %}
+{%- endfor %}
+
+
+{{ type_name }} = Annotated[
+    t.Union[
+         {%- for variant in type.variants -%}
+            {{ type_name }}{{ variant.name | to_upper_camel_case }}
+            {%- if not loop.last %}, {% endif %}
+        {%- endfor -%}
+    ],
+    Field(discriminator="{{ type.discriminator_field }}")
+]
+{%- if type.description is defined %}
+{{ type.description | to_doc_comment(style="python") }}
+{% endif %}

python/templates/types/struct_fields.py.jinja

@@ -0,0 +1,24 @@
+{% for field in type.fields %}
+    {%- if field.required and not field.nullable %}
+        {%- if field.name | to_lower_camel_case != field.name %}
+    {{ field.name | to_snake_case }}: {{ field.type.to_python() }} = Field(alias="{{ field.name }}")
+        {%- else %}
+    {{ field.name | to_snake_case }}: {{ field.type.to_python() }}
+        {%- endif %}
+    {%- elif field.required and field.nullable %}
+        {%- if field.name | to_lower_camel_case != field.name %}
+    {{ field.name | to_snake_case }}: t.Optional[{{ field.type.to_python() }}] = Field(alias="{{ field.name }}")
+        {%- else %}
+    {{ field.name | to_snake_case }}: t.Optional[{{ field.type.to_python() }}]
+        {%- endif %}
+    {%- else %}
+        {%- if field.name | to_lower_camel_case != field.name %}
+    {{ field.name | to_snake_case }}: t.Optional[{{ field.type.to_python() }}] = Field(default=None, alias="{{ field.name }}")
+        {%- else %}
+    {{ field.name | to_snake_case }}: t.Optional[{{ field.type.to_python() }}] = None
+        {%- endif %}
+    {%- endif %}
+    {%- if field.description is defined %}
+    {{ field.description | to_doc_comment(style="python") | indent(4) }}
+    {%- endif %}
+{% endfor %}