[ENUM] Default handler to use enum.name (#1252)

Jaseci-Labs · Sep 9, 2024 · 6719c96 · 6719c96
1 parent 6fa955b
commit 6719c96
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 13 deletions.
diff --git a/jac-cloud/jaclang_jaseci/core/architype.py b/jac-cloud/jaclang_jaseci/core/architype.py
@@ -1,13 +1,15 @@
 """Core constructs for Jac Language."""
 
-from dataclasses import asdict, dataclass, field, fields, is_dataclass
+from dataclasses import asdict as _asdict, dataclass, field, fields, is_dataclass
+from enum import Enum
 from os import getenv
 from pickle import dumps as pdumps
 from re import IGNORECASE, compile
 from typing import (
     Any,
     Callable,
     ClassVar,
+    Iterable,
     Mapping,
     TypeVar,
     cast,
@@ -53,6 +55,24 @@
 TBA = TypeVar("TBA", bound="BaseArchitype")
 
 
+def asdict_factory(data: Iterable[tuple]) -> dict[str, Any]:
+    """Parse dataclass to dict."""
+    _data = {}
+    for key, value in data:
+        if isinstance(value, Enum):
+            _data[key] = value.name
+        else:
+            _data[key] = value
+    return _data
+
+
+def asdict(obj: object) -> dict[str, Any]:
+    """Override dataclass asdict."""
+    if is_dataclass(obj) and not isinstance(obj, type):
+        return _asdict(obj, dict_factory=asdict_factory)
+    raise ValueError("Object is not a dataclass!")
+
+
 def architype_to_dataclass(cls: type[T], data: dict[str, Any], **kwargs: object) -> T:
     """Parse dict to architype."""
     _to_dataclass(cls, data)
@@ -89,6 +109,12 @@ def _to_dataclass(cls: type[T], data: dict[str, Any]) -> None:
                     ):
                         for key, value in enumerate(target):
                             target[key] = to_dataclass(inner_cls, value)
+                    elif (
+                        issubclass(hint, Enum)
+                        and isinstance(target, str)
+                        and (enum := hint.__members__.get(target))
+                    ):
+                        data[attr.name] = enum
 
 
 @dataclass

diff --git a/jac-cloud/jaclang_jaseci/core/context.py b/jac-cloud/jaclang_jaseci/core/context.py
@@ -1,7 +1,7 @@
 """Core constructs for Jac Language."""
 
 from contextvars import ContextVar
-from dataclasses import asdict, is_dataclass
+from dataclasses import is_dataclass
 from os import getenv
 from typing import Any, cast
 
@@ -19,6 +19,7 @@
     NodeAnchor,
     Permission,
     Root,
+    asdict,
 )
 from .memory import MongoDB
 

diff --git a/jac-cloud/jaclang_jaseci/tests/simple_graph.jac b/jac-cloud/jaclang_jaseci/tests/simple_graph.jac
@@ -38,6 +38,12 @@ class User:BaseUser: {
     }
 }
 
+enum Enum {
+    A = "a",
+    B = "b",
+    C = "c"
+}
+
 node A {
     has val: int;
 }
@@ -51,15 +57,15 @@ node C {
 }
 
 obj Child {
-    has val: int, arr: list, json: dict;
+    has val: int, arr: list, json: dict, enum_field: Enum;
 }
 
 obj Parent:Child: {
     has child: Child;
 }
 
 node Nested {
-    has val: int, arr: list, json: dict, parent: Parent;
+    has val: int, arr: list, json: dict, parent: Parent, enum_field: Enum;
 }
 
 walker create_graph {
@@ -145,9 +151,13 @@ walker create_nested_node {
                 child=Child(
                     val=2,
                     arr=[1, 2],
-                    json={"a": 1, "b": 2}
-                )
-            ));
+                    json={"a": 1, "b": 2},
+                    enum_field = Enum.C
+                ),
+                enum_field = Enum.B
+            ),
+            enum_field = Enum.A
+        );
         here ++> n;
         return n;
     }
@@ -159,25 +169,31 @@ walker update_nested_node {
         nested.parent.child.json["c"] = 3;
         nested.parent.child.arr.append(3);
         nested.parent.child.val = 3;
+        nested.parent.child.enum_field = Enum.A;
         nested.parent.json["b"] = 2;
         nested.parent.arr.append(2);
         nested.parent.val = 2;
+        nested.parent.enum_field = Enum.C;
         nested.json["a"] = 1;
         nested.arr.append(1);
         nested.val = 1;
+        nested.enum_field = Enum.B;
         return nested;
     }
 
     can enter_nested with Nested entry {
         here.parent.child.json["c"] = 3;
         here.parent.child.arr.append(3);
         here.parent.child.val = 3;
+        here.parent.child.enum_field = Enum.A;
         here.parent.json["b"] = 2;
         here.parent.arr.append(2);
         here.parent.val = 2;
+        here.parent.enum_field = Enum.C;
         here.json["a"] = 1;
         here.arr.append(1);
         here.val = 1;
+        here.enum_field = Enum.B;
         return here;
     }
 }
@@ -216,9 +232,13 @@ walker manual_create_nested_node {
                 child=Child(
                     val=2,
                     arr=[1, 2],
-                    json={"a": 1, "b": 2}
-                )
-            ));
+                    json={"a": 1, "b": 2},
+                    enum_field = Enum.C
+                ),
+                enum_field = Enum.B
+            ),
+            enum_field = Enum.A
+        );
         here ++> n;
         here.__jac__.apply();
 
@@ -237,12 +257,15 @@ walker manual_update_nested_node {
         nested.parent.child.json["c"] = 3;
         nested.parent.child.arr.append(3);
         nested.parent.child.val = 3;
+        nested.parent.child.enum_field = Enum.A;
         nested.parent.json["b"] = 2;
         nested.parent.arr.append(2);
         nested.parent.val = 2;
+        nested.parent.enum_field = Enum.C;
         nested.json["a"] = 1;
         nested.arr.append(1);
         nested.val = 1;
+        nested.enum_field = Enum.B;
         nested.__jac__.apply();
 
         # simulate no auto save

diff --git a/jac-cloud/jaclang_jaseci/tests/test_simple_graph.py b/jac-cloud/jaclang_jaseci/tests/test_simple_graph.py
@@ -223,8 +223,15 @@ def trigger_create_nested_node_test(self, manual: bool = False) -> None:
                     "val": 1,
                     "arr": [1],
                     "json": {"a": 1},
-                    "child": {"val": 2, "arr": [1, 2], "json": {"a": 1, "b": 2}},
+                    "child": {
+                        "val": 2,
+                        "arr": [1, 2],
+                        "json": {"a": 1, "b": 2},
+                        "enum_field": "C",
+                    },
+                    "enum_field": "B",
                 },
+                "enum_field": "A",
             },
             res["returns"][0]["context"],
         )
@@ -252,8 +259,11 @@ def trigger_update_nested_node_test(
                             "val": 3,
                             "arr": [1, 2, 3],
                             "json": {"a": 1, "b": 2, "c": 3},
+                            "enum_field": "A",
                         },
+                        "enum_field": "C",
                     },
+                    "enum_field": "B",
                 },
                 res["returns"][0]["context"],
             )
@@ -310,8 +320,15 @@ def trigger_access_validation_test(
                     "val": 1,
                     "arr": [1],
                     "json": {"a": 1},
-                    "child": {"val": 2, "arr": [1, 2], "json": {"a": 1, "b": 2}},
+                    "child": {
+                        "val": 2,
+                        "arr": [1, 2],
+                        "json": {"a": 1, "b": 2},
+                        "enum_field": "C",
+                    },
+                    "enum_field": "B",
                 },
+                "enum_field": "A",
             },
             nested_node["context"],
         )
@@ -354,8 +371,11 @@ def trigger_access_validation_test(
                         "val": 3,
                         "arr": [1, 2, 3],
                         "json": {"a": 1, "b": 2, "c": 3},
+                        "enum_field": "A",
                     },
+                    "enum_field": "C",
                 },
+                "enum_field": "B",
             },
             res["returns"][0]["context"],
         )
@@ -373,8 +393,15 @@ def trigger_access_validation_test(
                     "val": 1,
                     "arr": [1],
                     "json": {"a": 1},
-                    "child": {"val": 2, "arr": [1, 2], "json": {"a": 1, "b": 2}},
+                    "child": {
+                        "val": 2,
+                        "arr": [1, 2],
+                        "json": {"a": 1, "b": 2},
+                        "enum_field": "C",
+                    },
+                    "enum_field": "B",
                 },
+                "enum_field": "A",
             },
             res["returns"][0]["context"],
         )
@@ -409,8 +436,11 @@ def trigger_access_validation_test(
                         "val": 3,
                         "arr": [1, 2, 3],
                         "json": {"a": 1, "b": 2, "c": 3},
+                        "enum_field": "A",
                     },
+                    "enum_field": "C",
                 },
+                "enum_field": "B",
             },
             res["returns"][0]["context"],
         )
@@ -432,8 +462,11 @@ def trigger_access_validation_test(
                         "val": 3,
                         "arr": [1, 2, 3],
                         "json": {"a": 1, "b": 2, "c": 3},
+                        "enum_field": "A",
                     },
+                    "enum_field": "C",
                 },
+                "enum_field": "B",
             },
             res["returns"][0]["context"],
         )