feat: restore discriminated union types (#623)

Author: razor-x

codegen/layouts/partials/resource-dataclass.hbs

@@ -10,6 +10,17 @@
 {{#each nestedClasses}}
 {{> resource-dataclass isNested=true}}
 
+{{/each}}
+{{#each nestedUnions}}
+{{../memberIndent}}{{className}} = Union[{{#each variants}}{{className}}{{#unless @last}}, {{/unless}}{{/each}}]
+{{../memberIndent}}_{{className}}Variants = {
+{{#each variants}}
+{{#each values}}
+{{../../../memberIndent}}    {{pythonString this}}: {{../className}},
+{{/each}}
+{{/each}}
+{{../memberIndent}} }
+
 {{/each}}
 {{#each properties}}
 {{../memberIndent}}{{pythonIdentifier name}}: {{type}}
@@ -22,6 +33,6 @@
 {{/unless}}
 {{memberIndent}}    return cls(
 {{#each properties}}
-{{../memberIndent}}        {{pythonIdentifier name}}={{#if isObject}}cls.{{nestedClassName}}.from_dict(d.get("{{name}}")) if d.get("{{name}}") is not None else None{{else}}{{#if isObjectList}}[cls.{{nestedClassName}}.from_dict(i) for i in d.get("{{name}}") or []]{{else}}{{#if isDictParam}}DeepAttrDict({{/if}}d.get("{{name}}", None){{#if isDictParam}}){{/if}}{{/if}}{{/if}},
+{{../memberIndent}}        {{pythonIdentifier name}}={{#if isObject}}cls.{{nestedClassName}}.from_dict(d.get("{{name}}")) if d.get("{{name}}") is not None else None{{else}}{{#if isDiscriminatedObjectList}}[_from_discriminated_dict(i, cls._{{nestedClassName}}Variants, "{{discriminator}}") for i in d.get("{{name}}") or []]{{else}}{{#if isObjectList}}[cls.{{nestedClassName}}.from_dict(i) for i in d.get("{{name}}") or []]{{else}}{{#if isDictParam}}DeepAttrDict({{/if}}d.get("{{name}}", None){{#if isDictParam}}){{/if}}{{/if}}{{/if}}{{/if}},
 {{/each}}
 {{memberIndent}}    )

codegen/layouts/partials/route-method.hbs

@@ -24,16 +24,16 @@
 
         return {{#if isAsync}}await resolve_action_attempt_async{{else}}resolve_action_attempt{{/if}}(
             client=self.client,
-            action_attempt=ActionAttempt.from_dict(res["action_attempt"]),
+            action_attempt=action_attempt_from_dict(res["action_attempt"]),
             wait_for_action_attempt=wait_for_action_attempt
         )
 {{else if (eq returnType "None")}}
 
         return None
 {{else if (isListType returnType)}}
 
-        return [{{listItemType returnType}}.from_dict(item) for item in res{{#each returnPath}}["{{this}}"]{{/each}}]
+        return [{{fromDict (listItemType returnType)}}(item) for item in res{{#each returnPath}}["{{this}}"]{{/each}}]
 {{else}}
 
-        return {{returnType}}.from_dict(res{{#each returnPath}}["{{this}}"]{{/each}})
+        return {{fromDict returnType}}(res{{#each returnPath}}["{{this}}"]{{/each}})
 {{/if}}

codegen/layouts/resource.hbs

@@ -1,7 +1,42 @@
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union{{#if union}}, cast{{/if}}
 from dataclasses import dataclass
 from ..deep_attr_dict import DeepAttrDict
 from ..resource_mapping import ResourceMapping
 
+{{#if hasDiscriminatedLists}}
 
+def _from_discriminated_dict(
+    d: Any, variants: Dict[str, Any], discriminator: str
+) -> Any:
+    variant = variants.get(d.get(discriminator))
+    return DeepAttrDict(d) if variant is None else variant.from_dict(d)
+{{/if}}
+
+{{#each classes}}
 {{> resource-dataclass}}
+
+
+{{/each}}
+{{#if union}}
+{{union.className}} = Union[{{#each union.variants}}{{className}}{{#unless @last}}, {{/unless}}{{/each}}]
+
+{{union.variantsName}}: Dict[str, Any] = {
+{{#each union.variants}}
+{{#each values}}
+    {{pythonString this}}: {{../className}},
+{{/each}}
+{{/each}}
+}
+
+
+def {{union.fromDictName}}(d: Any) -> {{union.className}}:
+    """Deserialize a known {{union.discriminator}} variant.
+
+    Unknown discriminator values return ``DeepAttrDict`` so payloads from a
+    newer API remain readable. The static return type covers known variants.
+    """
+    variant = {{union.variantsName}}.get(d.get("{{union.discriminator}}"))
+    if variant is None:
+        return cast({{union.className}}, DeepAttrDict(d))
+    return variant.from_dict(d)
+{{/if}}

codegen/layouts/resources-index.hbs

@@ -1,3 +1,3 @@
 {{#each resources}}
-from .{{moduleName}} import {{className}}
+from .{{moduleName}} import ({{#each exports}}{{this}}{{#unless @last}}, {{/unless}}{{/each}})
 {{/each}}

codegen/lib/handlebars-helpers.ts

@@ -53,6 +53,8 @@ export const indent = (value: string, spaces: number): string =>
 export const pythonIdentifier = (name: string): string =>
   PYTHON_KEYWORDS.has(name) ? `${name}_` : name
 
+export const pythonString = (value: string): string => JSON.stringify(value)
+
 // A param the API documents as nullable may be set to the NULL sentinel, which
 // the client serializes to null. Params that are merely optional may not: they
 // are omitted by passing None, and sending null would unset a value instead.
@@ -62,3 +64,9 @@ export const nullableType = (type: string, isNullable: boolean): string =>
 export const isListType = (type: string): boolean => type.startsWith('List[')
 
 export const listItemType = (type: string): string => type.slice(5, -1)
+
+export const fromDict = (type: string): string => {
+  if (type === 'SeamEvent') return 'seam_event_from_dict'
+  if (type === 'ActionAttempt') return 'action_attempt_from_dict'
+  return `${type}.from_dict`
+}