[KED-2705] Show decorated function's source code on code panel (#493)

package/kedro_viz/api/apps.py

@@ -32,7 +32,7 @@
 from pathlib import Path
 
 from fastapi import FastAPI
-from fastapi.responses import JSONResponse, HTMLResponse
+from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from jinja2 import Environment, FileSystemLoader
 

package/kedro_viz/models/graph.py

@@ -36,6 +36,7 @@
 from dataclasses import InitVar, dataclass, field
 from enum import Enum
 from pathlib import Path
+from types import FunctionType
 from typing import Any, Dict, List, Optional, Set, Union, cast
 
 from kedro.io import AbstractDataSet
@@ -282,6 +283,18 @@ def __post_init__(self, kedro_obj: KedroNode):
         self.modular_pipelines = self._expand_namespaces(kedro_obj.namespace)
 
 
+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
+
+
 @dataclass
 class TaskNodeMetadata(GraphNodeMetadata):
     """Represent the metadata of a TaskNode"""
@@ -300,7 +313,9 @@ class TaskNodeMetadata(GraphNodeMetadata):
 
     def __post_init__(self, task_node: TaskNode):
         kedro_node = cast(KedroNode, task_node.kedro_obj)
-        self.code = inspect.getsource(kedro_node._func)
+        self.code = inspect.getsource(
+            _extract_wrapped_func(cast(FunctionType, kedro_node._func))
+        )
         code_full_path = Path(inspect.getfile(kedro_node._func)).expanduser().resolve()
         try:
             filepath = code_full_path.relative_to(Path.cwd().parent)

package/tests/test_models/test_graph/test_graph_nodes.py

@@ -59,6 +59,25 @@ def identity(x):
     return x
 
 
+def decorator(fun):
+    """
+    Not the best way to write decorator
+    but trying to stay faithful to the bug report here:
+    https://github.com/quantumblacklabs/kedro-viz/issues/484
+    """
+
+    def _new_fun(*args, **kwargs):
+        return fun(*args, **kwargs)
+
+    _new_fun.__name__ = fun.__name__
+    return _new_fun
+
+
+@decorator
+def decorated(x):
+    return x
+
+
 class TestGraphNodeCreation:
     @pytest.mark.parametrize(
         "namespace,expected_modular_pipelines",
@@ -239,6 +258,29 @@ def identity(x):
         )
         assert task_node_metadata.parameters == {}
 
+    def test_task_node_metadata_with_decorated_func(self):
+        kedro_node = node(
+            decorated,
+            inputs="x",
+            outputs="y",
+            name="identity_node",
+            tags={"tag"},
+            namespace="namespace",
+        )
+        task_node = GraphNode.create_task_node(kedro_node)
+        task_node_metadata = TaskNodeMetadata(task_node=task_node)
+        assert task_node_metadata.code == dedent(
+            """\
+            @decorator
+            def decorated(x):
+                return x
+            """
+        )
+        assert task_node_metadata.filepath == str(
+            Path(__file__).relative_to(Path.cwd().parent).expanduser()
+        )
+        assert task_node_metadata.parameters == {}
+
     def test_data_node_metadata(self):
         dataset = CSVDataSet(filepath="/tmp/dataset.csv")
         data_node = GraphNode.create_data_node(