From c8f4983e6b8f3ff0db900406d3e2e904f30c39d0 Mon Sep 17 00:00:00 2001 From: Jitendra Gundaniya Date: Wed, 1 May 2024 17:02:59 +0100 Subject: [PATCH] new approach for tags Signed-off-by: Jitendra Gundaniya --- .../repositories/modular_pipelines.py | 41 +++++++++---------- .../kedro_viz/services/modular_pipelines.py | 1 - .../test_modular_pipelines.py | 36 +++++----------- 3 files changed, 31 insertions(+), 47 deletions(-) diff --git a/package/kedro_viz/data_access/repositories/modular_pipelines.py b/package/kedro_viz/data_access/repositories/modular_pipelines.py index a8f4ea6496..1886f684be 100644 --- a/package/kedro_viz/data_access/repositories/modular_pipelines.py +++ b/package/kedro_viz/data_access/repositories/modular_pipelines.py @@ -1,7 +1,7 @@ """`kedro_viz.data_access.repositories.modular_pipelines` defines repository to centralise access to modular pipelines data.""" -from typing import Dict, Optional, Set, Union +from typing import Dict, Optional, Union from kedro_viz.constants import ROOT_MODULAR_PIPELINE_ID from kedro_viz.models.flowchart import ( @@ -163,9 +163,25 @@ def add_output(self, modular_pipeline_id: str, output_node: GraphNode): self.tree[modular_pipeline_id].external_outputs.add(output_node.id) def add_tags(self, modular_pipeline_id: str, node_tags: set): - if modular_pipeline_id in self.tree: - self.tree[modular_pipeline_id].tags |= node_tags - + """ + Add tags to a modular pipeline. + + Args: + modular_pipeline_id: ID of the modular pipeline to add the tags to. + node_tags: The tags to add to the modular pipeline. + + Example: + >>> modular_pipelines = ModularPipelinesRepository() + >>> node_tags = {"tag1", "tag2"} + >>> modular_pipelines.add_tags("data_science", node_tags) + >>> data_science_pipeline = modular_pipelines.get_or_create_modular_pipeline( + ... "data_science" + ... ) + >>> assert "tag1" in data_science_pipeline.tags + >>> assert "tag2" in data_science_pipeline.tags + """ + if modular_pipeline_id in self.tree: + self.tree[modular_pipeline_id].tags |= node_tags def add_child(self, modular_pipeline_id: str, child: ModularPipelineChild): """Add a child to a modular pipeline. @@ -215,8 +231,6 @@ def extract_from_node(self, node: GraphNode) -> Optional[str]: return None modular_pipeline = self.get_or_create_modular_pipeline(modular_pipeline_id) - # Inherit tags from the nodes of the modular pipeline. - self.inherit_tags_recursive(modular_pipeline_id, node.tags) # Add the node's registered pipelines to the modular pipeline's registered pipelines. # Basically this means if the node belongs to the "__default__" pipeline, for example, @@ -231,21 +245,6 @@ def extract_from_node(self, node: GraphNode) -> Optional[str]: ) return modular_pipeline_id - def inherit_tags_recursive(self, modular_pipeline_id: str, tags: Set[str]): - """Recursively collects a set of tags from a modular pipeline to all of its - child modular pipelines. - Args: - modular_pipeline_id: ID of the modular pipeline to check existence in the repository. - tags: A set of tags to be added to the modular pipeline and its children. - """ - - modular_pipeline = self.tree.get(modular_pipeline_id) - if modular_pipeline: - modular_pipeline.tags.update(tags) - for child in modular_pipeline.children: - if child.type == GraphNodeType.MODULAR_PIPELINE: - self.inherit_tags_recursive(child.id, modular_pipeline.tags) - def has_modular_pipeline(self, modular_pipeline_id: str) -> bool: """Return whether this modular pipeline repository has a given modular pipeline ID. Args: diff --git a/package/kedro_viz/services/modular_pipelines.py b/package/kedro_viz/services/modular_pipelines.py index 2487dca8de..bc568e99ff 100644 --- a/package/kedro_viz/services/modular_pipelines.py +++ b/package/kedro_viz/services/modular_pipelines.py @@ -98,5 +98,4 @@ def expand_tree( expanded_tree[parent_id].external_outputs.update( modular_pipeline_node.external_outputs ) - expanded_tree[parent_id].tags.update(modular_pipeline_node.tags) return expanded_tree diff --git a/package/tests/test_data_access/test_repositories/test_modular_pipelines.py b/package/tests/test_data_access/test_repositories/test_modular_pipelines.py index cdc72ef1b0..69c9a8bf81 100644 --- a/package/tests/test_data_access/test_repositories/test_modular_pipelines.py +++ b/package/tests/test_data_access/test_repositories/test_modular_pipelines.py @@ -36,31 +36,6 @@ def test_extract_from_node(self, identity): modular_pipelines.extract_from_node(task_node) assert modular_pipelines.has_modular_pipeline("data_science") - def test_tags_inheritance(self, identity): - task_node = GraphNode.create_task_node( - node( - identity, - inputs="x", - outputs=None, - namespace="parent", - tags={"tag1", "tag2"}, - ) - ) - - modular_pipelines = ModularPipelinesRepository() - modular_pipelines.add_child( - "parent", - ModularPipelineChild( - id="parent.child", type=GraphNodeType.MODULAR_PIPELINE - ), - ) - - modular_pipelines.extract_from_node(task_node) - modular_pipeline = modular_pipelines.get_or_create_modular_pipeline("parent") - - assert "tag1" in modular_pipeline.tags - assert "tag2" in modular_pipeline.tags - def test_add_input(self): kedro_dataset = CSVDataset(filepath="foo.csv") modular_pipelines = ModularPipelinesRepository() @@ -108,3 +83,14 @@ def test_add_output_should_raise_if_adding_non_data_node(self, identity): modular_pipelines = ModularPipelinesRepository() with pytest.raises(ValueError): modular_pipelines.add_output("data_science", task_node) + + def test_add_tags(self): + modular_pipelines = ModularPipelinesRepository() + node_tags = {"tag1", "tag2"} + modular_pipelines.get_or_create_modular_pipeline("data_science") + modular_pipelines.add_tags("data_science", node_tags) + data_science_pipeline = modular_pipelines.get_or_create_modular_pipeline( + "data_science" + ) + assert "tag1" in data_science_pipeline.tags + assert "tag2" in data_science_pipeline.tags