From dict to model_dump

pier-digital · Oct 10, 2023 · 2288394 · 2288394
1 parent 6cb6884
commit 2288394
Show file tree

Hide file tree

Showing 14 changed files with 25 additions and 20 deletions.
diff --git a/retrack/engine/parser.py b/retrack/engine/parser.py
@@ -166,7 +166,7 @@ def _set_execution_order(self):
     def get_node_connections(
         self, node_id: str, is_input: bool = True, filter_by_connector=None
     ):
-        node_dict = self.get_by_id(node_id).dict(by_alias=True)
+        node_dict = self.get_by_id(node_id).model_dump(by_alias=True)
 
         connectors = node_dict.get("inputs" if is_input else "outputs", {})
         result = []

diff --git a/retrack/engine/runner.py b/retrack/engine/runner.py
@@ -111,7 +111,7 @@ def _create_initial_state_from_payload(
     ) -> pd.DataFrame:
         """Create initial state from payload. This is the first step of the runner."""
         validated_payload = self.request_manager.validate(payload)
-        validated_payload = pd.DataFrame([p.dict() for p in validated_payload])
+        validated_payload = pd.DataFrame([p.model_dump() for p in validated_payload])
 
         state_df = pd.DataFrame([])
         for node_id, input_name in self.input_columns.items():
@@ -169,7 +169,7 @@ def __run_node(self, node_id: str):
             return
 
         input_params = self.__get_input_params(
-            node_id, node.dict(by_alias=True), current_node_filter
+            node_id, node.model_dump(by_alias=True), current_node_filter
         )
         output = node.run(**input_params)
 

diff --git a/retrack/nodes/base.py b/retrack/nodes/base.py
@@ -34,14 +34,18 @@ class NodeMemoryType(str, enum.Enum):
 # Connection Models
 ###############################################################
 
+def cast_int_to_str(v: typing.Any, info: pydantic.ValidationInfo) -> str:
+    return str(v)
+
+CastedToStringType = typing.Annotated[typing.Any, pydantic.BeforeValidator(cast_int_to_str)]
 
 class OutputConnectionItemModel(pydantic.BaseModel):
-    node: str
+    node: CastedToStringType
     input_: str = pydantic.Field(alias="input")
 
 
 class InputConnectionItemModel(pydantic.BaseModel):
-    node: str
+    node: CastedToStringType
     output: str = pydantic.Field(alias="output")
 
 
@@ -57,11 +61,12 @@ class InputConnectionModel(pydantic.BaseModel):
 # Base Node
 ###############################################################
 
-
 class BaseNode(pydantic.BaseModel):
-    id: str
-    inputs: typing.Optional[typing.Dict[str, InputConnectionModel]] = None
-    outputs: typing.Optional[typing.Dict[str, OutputConnectionModel]] = None
+    id: CastedToStringType
+    inputs: typing.Optional[typing.Dict[CastedToStringType, InputConnectionModel]] = None
+    outputs: typing.Optional[typing.Dict[CastedToStringType, OutputConnectionModel]] = None
+
+
 
     def run(self, **kwargs) -> typing.Dict[str, typing.Any]:
         return {}

diff --git a/tests/test_nodes/test_check.py b/tests/test_nodes/test_check.py
@@ -25,7 +25,7 @@ def test_check_node(node_schema):
     assert isinstance(check_node, Check)
     assert isinstance(check_node.data.operator, CheckOperator)
 
-    assert check_node.dict(by_alias=True) == {
+    assert check_node.model_dump(by_alias=True) == {
         "id": "8",
         "data": {"operator": CheckOperator.EQUAL},
         "inputs": {

diff --git a/tests/test_nodes/test_contains.py b/tests/test_nodes/test_contains.py
@@ -17,7 +17,7 @@ def test_Contains_node():
 
     assert isinstance(Contains_node, Contains)
 
-    assert Contains_node.dict(by_alias=True) == {
+    assert Contains_node.model_dump(by_alias=True) == {
         "id": "9",
         "inputs": {
             "input_list": {"connections": []},

diff --git a/tests/test_nodes/test_datetime.py b/tests/test_nodes/test_datetime.py
@@ -22,7 +22,7 @@ def test_current_year_node(current_year_input_data):
 
     assert isinstance(current_year_node, CurrentYear)
 
-    assert current_year_node.dict(by_alias=True) == {
+    assert current_year_node.model_dump(by_alias=True) == {
         "id": "18",
         "inputs": {
             "input_void": {"connections": []},

diff --git a/tests/test_nodes/test_endswith.py b/tests/test_nodes/test_endswith.py
@@ -20,7 +20,7 @@ def test_EndsWith_node():
 
     assert isinstance(EndsWith_node, EndsWith)
 
-    assert EndsWith_node.dict(by_alias=True) == {
+    assert EndsWith_node.model_dump(by_alias=True) == {
         "id": "11",
         "inputs": {
             "input_value_0": {"connections": []},

diff --git a/tests/test_nodes/test_endswithany.py b/tests/test_nodes/test_endswithany.py
@@ -17,7 +17,7 @@ def test_EndsWithAny_node():
 
     assert isinstance(EndsWithAny_node, EndsWithAny)
 
-    assert EndsWithAny_node.dict(by_alias=True) == {
+    assert EndsWithAny_node.model_dump(by_alias=True) == {
         "id": "13",
         "inputs": {
             "input_value": {"connections": []},

diff --git a/tests/test_nodes/test_input.py b/tests/test_nodes/test_input.py
@@ -6,4 +6,4 @@ def test_input_node(
 ):
     input_node = Input(**valid_input_dict_before_validation)
 
-    assert input_node.dict(by_alias=True) == valid_input_dict_after_validation
+    assert input_node.model_dump(by_alias=True) == valid_input_dict_after_validation
diff --git a/tests/test_nodes/test_logic.py b/tests/test_nodes/test_logic.py
@@ -45,7 +45,7 @@ def test_Logic_node():
     assert isinstance(Or_node, Or)
     assert isinstance(Not_node, Not)
 
-    assert And_node.dict(by_alias=True) == {
+    assert And_node.model_dump(by_alias=True) == {
         "id": "15",
         "inputs": {
             "input_bool_0": {"connections": []},

diff --git a/tests/test_nodes/test_math.py b/tests/test_nodes/test_math.py
@@ -39,7 +39,7 @@ def test_math_node(math_operator_input_data):
     assert isinstance(math_node, Math)
     assert isinstance(math_node.data.operator, MathOperator)
 
-    assert math_node.dict(by_alias=True) == {
+    assert math_node.model_dump(by_alias=True) == {
         "id": "18",
         "data": {"operator": MathOperator.SUM},
         "inputs": {

diff --git a/tests/test_nodes/test_start.py b/tests/test_nodes/test_start.py
@@ -24,7 +24,7 @@ def test_start_node():
 
     start_node = Start(**input_data)
 
-    assert start_node.dict(by_alias=True) == {
+    assert start_node.model_dump(by_alias=True) == {
         "id": "0",
         "outputs": {
             "output_up_void": {

diff --git a/tests/test_nodes/test_startswith.py b/tests/test_nodes/test_startswith.py
@@ -20,7 +20,7 @@ def test_StartsWith_node():
 
     assert isinstance(StartsWith_node, StartsWith)
 
-    assert StartsWith_node.dict(by_alias=True) == {
+    assert StartsWith_node.model_dump(by_alias=True) == {
         "id": "10",
         "inputs": {
             "input_value_0": {"connections": []},

diff --git a/tests/test_nodes/test_startswithany.py b/tests/test_nodes/test_startswithany.py
@@ -17,7 +17,7 @@ def test_StartsWithAny_node():
 
     assert isinstance(StartsWithAny_node, StartsWithAny)
 
-    assert StartsWithAny_node.dict(by_alias=True) == {
+    assert StartsWithAny_node.model_dump(by_alias=True) == {
         "id": "13",
         "inputs": {
             "input_value": {"connections": []},