feat: add python discriminator support

Author: 01Joseph-Hwang10

jtd_codebuild/generators/python/_generator.py

@@ -21,18 +21,32 @@ def generate(self, target: TypescriptTarget) -> None:
         propertyFormat = target.propertyFormat
 
         if propertyFormat:
-            convert_schema_definition_keys = compose_left(
+            convert_property_keys = compose_left(
                 curry(convert_keys_case)(propertyFormat),
                 dict,
             )
+
+            def convert_definition_keys_case(definition: dict) -> None:
+                if "discriminator" in definition:
+                    definition["discriminator"] = caseconverter(
+                        propertyFormat,
+                        definition["discriminator"],
+                    )
+                    for case_ in definition["mapping"].values():
+                        convert_definition_keys_case(case_)
+                else:
+                    properties = definition.get("properties", None)
+                    if properties:
+                        definition["properties"] = convert_property_keys(properties)
+                    optionalProperties = definition.get("optionalProperties", None)
+                    if optionalProperties:
+                        definition["optionalProperties"] = convert_property_keys(
+                            optionalProperties
+                        )
+
             for class_name in definitions.keys():
                 definition = definitions[class_name]
-                definition["properties"] = convert_schema_definition_keys(
-                    definition.get("properties", {})
-                )
-                definition["optionalProperties"] = convert_schema_definition_keys(
-                    definition.get("optionalProperties", {})
-                )
+                convert_definition_keys_case(definition)
 
         write_json(jtd_schema_path, jtd_schema)
 

jtd_codebuild/generators/typescript/_generator.py

@@ -0,0 +1,3 @@
+# flake8: noqa: F401
+
+from ._find import find, find_one

jtd_codebuild/utils/list/__init__.py

@@ -0,0 +1,26 @@
+from typing import TypeVar, Callable
+from toolz import pipe
+from toolz.curried import filter
+
+T = TypeVar("T")
+
+
+def find(
+    predicate: Callable[[T], bool],
+    nodes: list[T],
+) -> list[T]:
+    return pipe(
+        nodes,
+        filter(predicate),
+        list,
+    )
+
+
+def find_one(
+    predicate: Callable[[T], bool],
+    nodes: list[T],
+) -> T | None:
+    found = find(predicate, nodes)
+    if len(found) > 0:
+        return found[0]
+    return None

jtd_codebuild/utils/list/_find.py

@@ -0,0 +1,21 @@
+Event:
+  discriminator: eventType
+  mapping:
+    USER_CREATED:
+      properties:
+        id:
+          type: string
+    USER_PAYMENT_PLAN_CHANGED:
+      properties:
+        id:
+          type: string
+        plan:
+          enum:
+            - Free
+            - Paid
+    USER_DELETED:
+      properties:
+        id:
+          type: string
+        softDelete:
+          type: boolean

tests/python_targets/project/src/event.jtd.yml

@@ -10,3 +10,10 @@ Person:
 
       metadata:
         description: The name of the person.
+
+    status:
+      enum:
+        - happy
+        - sad
+        - angry
+        - confused

tests/python_targets/project/src/person.jtd.yml

@@ -0,0 +1,21 @@
+Event:
+  discriminator: eventType
+  mapping:
+    USER_CREATED:
+      properties:
+        id:
+          type: string
+    USER_PAYMENT_PLAN_CHANGED:
+      properties:
+        id:
+          type: string
+        plan:
+          enum:
+            - Free
+            - Paid
+    USER_DELETED:
+      properties:
+        id:
+          type: string
+        softDelete:
+          type: boolean

tests/typescript_targets/project/src/event.jtd.yml

@@ -10,3 +10,10 @@ Person:
 
       metadata:
         description: The name of the person.
+
+    status:
+      enum:
+        - happy
+        - sad
+        - angry
+        - confused