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
  )
  )