Breakdown flowchart models into separate files (#2144)

* imitial

Signed-off-by: Sajid Alam <[email protected]>

* update

Signed-off-by: Sajid Alam <[email protected]>

* split into modular pipelines

Signed-off-by: Sajid Alam <[email protected]>

* remove comment

Signed-off-by: Sajid Alam <[email protected]>

* move GraphNodeType to nodes

Signed-off-by: Sajid Alam <[email protected]>

* refactor

Signed-off-by: Sajid Alam <[email protected]>

* fix refactors

Signed-off-by: Sajid Alam <[email protected]>

* fix imports

Signed-off-by: Sajid Alam <[email protected]>

* resolve circular dependency

Signed-off-by: Sajid Alam <[email protected]>

* fix tests

Signed-off-by: Sajid Alam <[email protected]>

* lint

Signed-off-by: Sajid Alam <[email protected]>

* changes based on review

Signed-off-by: Sajid Alam <[email protected]>

* split flowchart test file

Signed-off-by: Sajid Alam <[email protected]>

* Update node_metadata.py

Signed-off-by: Sajid Alam <[email protected]>

* Update ruff.toml

Signed-off-by: Sajid Alam <[email protected]>

* lint

Signed-off-by: Sajid Alam <[email protected]>

* move test files

Signed-off-by: Sajid Alam <[email protected]>

* moved to named_entities.py

Signed-off-by: Sajid Alam <[email protected]>

---------

Signed-off-by: Sajid Alam <[email protected]>

package/kedro_viz/api/rest/responses.py

@@ -12,15 +12,13 @@
 
 from kedro_viz.api.rest.utils import get_package_compatibilities
 from kedro_viz.data_access import data_access_manager
-from kedro_viz.models.flowchart import (
-    DataNode,
+from kedro_viz.models.flowchart.node_metadata import (
     DataNodeMetadata,
     ParametersNodeMetadata,
-    TaskNode,
     TaskNodeMetadata,
-    TranscodedDataNode,
     TranscodedDataNodeMetadata,
 )
+from kedro_viz.models.flowchart.nodes import DataNode, TaskNode, TranscodedDataNode
 from kedro_viz.models.metadata import Metadata, PackageCompatibility
 
 logger = logging.getLogger(__name__)

package/kedro_viz/data_access/managers.py

@@ -20,15 +20,15 @@
 
 from kedro_viz.constants import DEFAULT_REGISTERED_PIPELINE_ID, ROOT_MODULAR_PIPELINE_ID
 from kedro_viz.integrations.utils import UnavailableDataset
-from kedro_viz.models.flowchart import (
+from kedro_viz.models.flowchart.edge import GraphEdge
+from kedro_viz.models.flowchart.model_utils import GraphNodeType
+from kedro_viz.models.flowchart.named_entities import RegisteredPipeline
+from kedro_viz.models.flowchart.nodes import (
     DataNode,
-    GraphEdge,
     GraphNode,
-    GraphNodeType,
     ModularPipelineChild,
     ModularPipelineNode,
     ParametersNode,
-    RegisteredPipeline,
     TaskNode,
     TranscodedDataNode,
 )

package/kedro_viz/data_access/repositories/graph.py

@@ -3,7 +3,8 @@
 
 from typing import Dict, Generator, List, Optional, Set
 
-from kedro_viz.models.flowchart import GraphEdge, GraphNode
+from kedro_viz.models.flowchart.edge import GraphEdge
+from kedro_viz.models.flowchart.nodes import GraphNode
 
 
 class GraphNodesRepository:

package/kedro_viz/data_access/repositories/modular_pipelines.py

@@ -8,9 +8,9 @@
 from kedro.pipeline.node import Node as KedroNode
 
 from kedro_viz.constants import ROOT_MODULAR_PIPELINE_ID
-from kedro_viz.models.flowchart import (
+from kedro_viz.models.flowchart.model_utils import GraphNodeType
+from kedro_viz.models.flowchart.nodes import (
     GraphNode,
-    GraphNodeType,
     ModularPipelineChild,
     ModularPipelineNode,
 )

package/kedro_viz/data_access/repositories/registered_pipelines.py

@@ -4,7 +4,7 @@
 from collections import OrderedDict, defaultdict
 from typing import Dict, List, Optional, Set
 
-from kedro_viz.models.flowchart import RegisteredPipeline
+from kedro_viz.models.flowchart.named_entities import RegisteredPipeline
 
 
 class RegisteredPipelinesRepository:

package/kedro_viz/data_access/repositories/tags.py

@@ -3,7 +3,7 @@
 
 from typing import Iterable, List, Set
 
-from kedro_viz.models.flowchart import Tag
+from kedro_viz.models.flowchart.named_entities import Tag
 
 
 class TagsRepository:

package/kedro_viz/models/flowchart/edge.py

@@ -0,0 +1,15 @@
+"""`kedro_viz.models.flowchart.edge` defines data models to represent Kedro edges in a viz graph."""
+
+from pydantic import BaseModel
+
+
+class GraphEdge(BaseModel, frozen=True):
+    """Represent an edge in the graph
+
+    Args:
+        source (str): The id of the source node.
+        target (str): The id of the target node.
+    """
+
+    source: str
+    target: str

package/kedro_viz/models/flowchart/model_utils.py

@@ -0,0 +1,45 @@
+"""`kedro_viz.models.flowchart.model_utils` defines utils for Kedro entities in a viz graph."""
+
+import logging
+from enum import Enum
+from types import FunctionType
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_filepath(dataset_description: Dict[str, Any]) -> Optional[str]:
+    """
+    Extract the file path from a dataset description dictionary.
+    """
+    filepath = dataset_description.get("filepath") or dataset_description.get("path")
+    return str(filepath) if filepath else None
+
+
+def _extract_wrapped_func(func: FunctionType) -> FunctionType:
+    """Extract a wrapped decorated function to inspect the source code if available.
+    Adapted from https://stackoverflow.com/a/43506509/1684058
+    """
+    if func.__closure__ is None:
+        return func
+    closure = (c.cell_contents for c in func.__closure__)
+    wrapped_func = next((c for c in closure if isinstance(c, FunctionType)), None)
+    # return the original function if it's not a decorated function
+    return func if wrapped_func is None else wrapped_func
+
+
+# =============================================================================
+# Shared base classes and enumerations for model components
+# =============================================================================
+
+
+class GraphNodeType(str, Enum):
+    """Represent all possible node types in the graph representation of a Kedro pipeline.
+    The type needs to inherit from str as well so FastAPI can serialise it. See:
+    https://fastapi.tiangolo.com/tutorial/path-params/#working-with-python-enumerations
+    """
+
+    TASK = "task"
+    DATA = "data"
+    PARAMETERS = "parameters"
+    MODULAR_PIPELINE = "modularPipeline"  # CamelCase for frontend compatibility

package/kedro_viz/models/flowchart/named_entities.py

@@ -0,0 +1,41 @@
+"""kedro_viz.models.flowchart.named_entities` defines data models for representing named entities
+such as tags and registered pipelines within a Kedro visualization graph."""
+
+from typing import Optional
+
+from pydantic import BaseModel, Field, ValidationInfo, field_validator
+
+
+class NamedEntity(BaseModel):
+    """Represent a named entity (Tag/Registered Pipeline) in a Kedro project
+    Args:
+        id (str): Id of the registered pipeline
+
+    Raises:
+        AssertionError: If id is not supplied during instantiation
+    """
+
+    id: str
+    name: Optional[str] = Field(
+        default=None,
+        validate_default=True,
+        description="The name of the entity",
+    )
+
+    @field_validator("name")
+    @classmethod
+    def set_name(cls, _, info: ValidationInfo):
+        """Ensures that the 'name' field is set to the value of 'id' if 'name' is not provided."""
+        assert "id" in info.data
+        return info.data["id"]
+
+
+class RegisteredPipeline(NamedEntity):
+    """Represent a registered pipeline in a Kedro project."""
+
+
+class Tag(NamedEntity):
+    """Represent a tag in a Kedro project."""
+
+    def __hash__(self) -> int:
+        return hash(self.id)