From c580f3138bb96142f2e5a847d0cbb986f7e96333 Mon Sep 17 00:00:00 2001 From: Jonas Platte Date: Fri, 28 Feb 2025 13:27:15 +0100 Subject: [PATCH] python: Update codegen templates to support struct enums --- python/templates/component_type.py.jinja | 10 +++-- python/templates/types/struct.py.jinja | 27 +------------ python/templates/types/struct_enum.py.jinja | 39 +++++++++++++++++++ python/templates/types/struct_fields.py.jinja | 24 ++++++++++++ 4 files changed, 71 insertions(+), 29 deletions(-) create mode 100644 python/templates/types/struct_enum.py.jinja create mode 100644 python/templates/types/struct_fields.py.jinja diff --git a/python/templates/component_type.py.jinja b/python/templates/component_type.py.jinja index 387314604..3e0f77872 100644 --- a/python/templates/component_type.py.jinja +++ b/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 %} diff --git a/python/templates/types/struct.py.jinja b/python/templates/types/struct.py.jinja index b7cc463b5..18f58b850 100644 --- a/python/templates/types/struct.py.jinja +++ b/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" %} diff --git a/python/templates/types/struct_enum.py.jinja b/python/templates/types/struct_enum.py.jinja new file mode 100644 index 000000000..019481d6f --- /dev/null +++ b/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 %} diff --git a/python/templates/types/struct_fields.py.jinja b/python/templates/types/struct_fields.py.jinja new file mode 100644 index 000000000..52b76d50f --- /dev/null +++ b/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 %}