Refactor Flowchart models from dataclass to pydantic base models (#1565)

* create pydantic flowchart classes Signed-off-by: ravi-kumar-pilla <[email protected]> * fixing pydantic class conversions Signed-off-by: ravi-kumar-pilla <[email protected]> * minor modifications and code comments Signed-off-by: ravi-kumar-pilla <[email protected]> * fix pytests for flowchart and managers Signed-off-by: ravi-kumar-pilla <[email protected]> * fix pytests Signed-off-by: ravi-kumar-pilla <[email protected]> * fix lint errors and pytests Signed-off-by: ravi-kumar-pilla <[email protected]> * fix lint issues and merge main Signed-off-by: ravi-kumar-pilla <[email protected]> * fix initialization issues and lint issues Signed-off-by: ravi-kumar-pilla <[email protected]> * fix push issue Signed-off-by: ravi-kumar-pilla <[email protected]> * revert class method to static method Signed-off-by: ravi-kumar-pilla <[email protected]> * update method comments Signed-off-by: ravi-kumar-pilla <[email protected]> * revert back method shuffles for better PR reviews Signed-off-by: ravi-kumar-pilla <[email protected]> * addressing PR comments Signed-off-by: ravi-kumar-pilla <[email protected]> * not-working version of pydantic shift Signed-off-by: ravi-kumar-pilla <[email protected]> * non-working p2 Signed-off-by: ravi-kumar-pilla <[email protected]> * not-working v3 Signed-off-by: ravi-kumar-pilla <[email protected]> * modify metadata classes and update pytests Signed-off-by: ravi-kumar-pilla <[email protected]> * fix all pytests Signed-off-by: ravi-kumar-pilla <[email protected]> * fix pytest for coverage Signed-off-by: ravi-kumar-pilla <[email protected]> * address PR comments1 Signed-off-by: ravi-kumar-pilla <[email protected]> * fix lint issues Signed-off-by: ravi-kumar-pilla <[email protected]> * add hash for pylint Signed-off-by: ravi-kumar-pilla <[email protected]> * fix lint errors Signed-off-by: ravi-kumar-pilla <[email protected]> * fix lint errors Signed-off-by: ravi-kumar-pilla <[email protected]> * create base class for tag and registered pipeline Signed-off-by: ravi-kumar-pilla <[email protected]> * add release note Signed-off-by: ravi-kumar-pilla <[email protected]> --------- Signed-off-by: ravi-kumar-pilla <[email protected]> Co-authored-by: Tynan DeBold <[email protected]>
kedro-org · Nov 14, 2023 · 2eb9c0a · 2eb9c0a
1 parent a38770b
commit 2eb9c0a
Show file tree

Hide file tree

Showing 15 changed files with 755 additions and 405 deletions.
diff --git a/RELEASE.md b/RELEASE.md
@@ -13,6 +13,7 @@ Please follow the established format:
 - Fix dataset factory patterns in Experiment Tracking. (#1588)
 - Improved feedback for copy to clipboard feature. (#1614)
 - Ensure Kedro-Viz works when hosted on a URL subpath. (#1621)
+- Refactor flowchart dataclasses to pydantic base models. (#1565)
 
 # Release 6.6.1
 

diff --git a/package/.pylintrc b/package/.pylintrc
@@ -22,7 +22,7 @@ jobs=1
 
 # List of plugins (as comma separated values of python modules names) to load,
 # usually to register additional checkers.
-load-plugins=pylint.extensions.docparams
+load-plugins=pylint.extensions.docparams,pylint_pydantic
 
 # Pickle collected data for later comparisons.
 persistent=yes

diff --git a/package/kedro_viz/api/graphql/serializers.py b/package/kedro_viz/api/graphql/serializers.py
@@ -10,10 +10,9 @@
 
 from strawberry import ID
 
+from kedro_viz.api.graphql.types import Run
 from kedro_viz.models.experiment_tracking import RunModel, UserRunDetailsModel
 
-from .types import Run
-
 
 def format_run(
     run_id: str, run_blob: Dict, user_run_details: Optional[UserRunDetailsModel] = None

diff --git a/package/kedro_viz/api/rest/responses.py b/package/kedro_viz/api/rest/responses.py
@@ -1,5 +1,5 @@
 """`kedro_viz.api.rest.responses` defines REST response types."""
-# pylint: disable=missing-class-docstring,too-few-public-methods,invalid-name
+# pylint: disable=missing-class-docstring,invalid-name
 import abc
 import logging
 from typing import Any, Dict, List, Optional, Union
@@ -310,18 +310,18 @@ def get_default_response() -> GraphAPIResponse:
     )
 
     return GraphAPIResponse(
-        nodes=data_access_manager.get_nodes_for_registered_pipeline(  # type: ignore
+        nodes=data_access_manager.get_nodes_for_registered_pipeline(
             default_selected_pipeline_id
         ),
-        edges=data_access_manager.get_edges_for_registered_pipeline(  # type: ignore
+        edges=data_access_manager.get_edges_for_registered_pipeline(
             default_selected_pipeline_id
         ),
         tags=data_access_manager.tags.as_list(),
         layers=data_access_manager.get_sorted_layers_for_registered_pipeline(
             default_selected_pipeline_id
         ),
         pipelines=data_access_manager.registered_pipelines.as_list(),
-        modular_pipelines=modular_pipelines_tree,  # type: ignore
+        modular_pipelines=modular_pipelines_tree,
         selected_pipeline=default_selected_pipeline_id,
     )
 
@@ -336,15 +336,15 @@ def get_node_metadata_response(node_id: str):
         return JSONResponse(content={})
 
     if isinstance(node, TaskNode):
-        return TaskNodeMetadata(node)
+        return TaskNodeMetadata(task_node=node)
 
     if isinstance(node, DataNode):
-        return DataNodeMetadata(node)
+        return DataNodeMetadata(data_node=node)
 
     if isinstance(node, TranscodedDataNode):
-        return TranscodedDataNodeMetadata(node)
+        return TranscodedDataNodeMetadata(transcoded_data_node=node)
 
-    return ParametersNodeMetadata(node)
+    return ParametersNodeMetadata(parameters_node=node)
 
 
 def get_selected_pipeline_response(registered_pipeline_id: str):
@@ -361,10 +361,10 @@ def get_selected_pipeline_response(registered_pipeline_id: str):
     )
 
     return GraphAPIResponse(
-        nodes=data_access_manager.get_nodes_for_registered_pipeline(  # type: ignore
+        nodes=data_access_manager.get_nodes_for_registered_pipeline(
             registered_pipeline_id
         ),
-        edges=data_access_manager.get_edges_for_registered_pipeline(  # type: ignore
+        edges=data_access_manager.get_edges_for_registered_pipeline(
             registered_pipeline_id
         ),
         tags=data_access_manager.tags.as_list(),
@@ -373,7 +373,7 @@ def get_selected_pipeline_response(registered_pipeline_id: str):
         ),
         pipelines=data_access_manager.registered_pipelines.as_list(),
         selected_pipeline=registered_pipeline_id,
-        modular_pipelines=modular_pipelines_tree,  # type: ignore
+        modular_pipelines=modular_pipelines_tree,
     )
 
 

diff --git a/package/kedro_viz/data_access/managers.py b/package/kedro_viz/data_access/managers.py
@@ -148,7 +148,7 @@ def add_pipeline(self, registered_pipeline_id: str, pipeline: KedroPipeline):
             task_node = self.add_node(registered_pipeline_id, node)
             self.registered_pipelines.add_node(registered_pipeline_id, task_node.id)
 
-            current_modular_pipeline = modular_pipelines.extract_from_node(task_node)
+            current_modular_pipeline_id = modular_pipelines.extract_from_node(task_node)
 
             # Add node's inputs as DataNode to the graph
             for input_ in node.inputs:
@@ -169,8 +169,8 @@ def add_pipeline(self, registered_pipeline_id: str, pipeline: KedroPipeline):
                 # The method `add_input` will take care of figuring out whether
                 # it is an internal or external input of the modular pipeline.
                 modular_pipelines.extract_from_node(input_node)
-                if current_modular_pipeline is not None:
-                    modular_pipelines.add_input(current_modular_pipeline, input_node)
+                if current_modular_pipeline_id is not None:
+                    modular_pipelines.add_input(current_modular_pipeline_id, input_node)
 
             # Add node outputs as DataNode to the graph.
             # It follows similar logic to adding inputs.
@@ -186,8 +186,10 @@ def add_pipeline(self, registered_pipeline_id: str, pipeline: KedroPipeline):
                     output_node.original_version = self.catalog.get_dataset(output)
 
                 modular_pipelines.extract_from_node(output_node)
-                if current_modular_pipeline is not None:
-                    modular_pipelines.add_output(current_modular_pipeline, output_node)
+                if current_modular_pipeline_id is not None:
+                    modular_pipelines.add_output(
+                        current_modular_pipeline_id, output_node
+                    )
 
     def add_node(self, registered_pipeline_id: str, node: KedroNode) -> TaskNode:
         """Add a Kedro node as a TaskNode to the NodesRepository
@@ -476,7 +478,9 @@ def create_modular_pipelines_tree_for_registered_pipeline(
             bad_inputs = modular_pipeline.inputs.intersection(descendants)
             for bad_input in bad_inputs:
                 digraph.remove_edge(bad_input, modular_pipeline_id)
-                edges.remove_edge(GraphEdge(bad_input, modular_pipeline_id))
+                edges.remove_edge(
+                    GraphEdge(source=bad_input, target=modular_pipeline_id)
+                )
                 node_dependencies[bad_input].remove(modular_pipeline_id)
 
         for node_id, node in self.nodes.as_dict().items():
@@ -488,7 +492,7 @@ def create_modular_pipelines_tree_for_registered_pipeline(
             if not node.modular_pipelines or node_id in root_parameters:
                 modular_pipelines_tree[ROOT_MODULAR_PIPELINE_ID].children.add(
                     ModularPipelineChild(
-                        node_id, self.nodes.get_node_by_id(node_id).type
+                        id=node_id, type=self.nodes.get_node_by_id(node_id).type
                     )
                 )