From cf959384cc6b3c77ef67e345efc69d3024c5c1cc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Y=C3=BBki=20VACHOT?= <yuki@vachot.fr>
Date: Wed, 17 Jan 2024 17:33:24 +0100
Subject: [PATCH] PipeGraph Decorator

---
 assets/pipegraph/pipegraph.py        | 12 ++++++------
 assets/pipegraph/pycallgraph_test.py |  6 ++++++
 2 files changed, 12 insertions(+), 6 deletions(-)
 create mode 100644 assets/pipegraph/pycallgraph_test.py

diff --git a/assets/pipegraph/pipegraph.py b/assets/pipegraph/pipegraph.py
index a261a41..ee92345 100644
--- a/assets/pipegraph/pipegraph.py
+++ b/assets/pipegraph/pipegraph.py
@@ -1,5 +1,7 @@
 import json
 import inspect
+import stack_data
+import uuid
 
 
 def get_node_type_from_func(func: object) -> str:
@@ -34,14 +36,10 @@ class PipeGraph(object):
     def __call__(self, *args, **kwargs):
         print(self.__name)
         print(self.__func.__doc__)
-        test = inspect.stack()
-        print(get_stack_frame_id(test[1][0].stack[0]))
+        print("-----------")
         print(self.add_node(
             doc=self.__func.__doc__,
         ))
-        print(self.add_link(
-            get_stack_frame_id()
-        ))
 
         return self.__func(*args, **kwargs)
 
@@ -52,7 +50,9 @@ class PipeGraph(object):
     def add_node(self, **kwargs):
         node_type = get_node_type_from_func(self.__func)
         current_id = PipeGraph.__node_id
-        node = {'id': current_id, 'name': self.__name, 'type': node_type}
+        frame = inspect.currentframe().f_back.f_back
+        current_co_name = stack_data.FrameInfo(frame).code.co_name
+        node = {'id': current_id, 'uuid': str(uuid.uuid4()), 'name': self.__name, 'type': node_type, 'before': current_co_name}
         for k, v in kwargs.items():
             node[k] = v
         PipeGraph.__json['nodes'].append(node)
diff --git a/assets/pipegraph/pycallgraph_test.py b/assets/pipegraph/pycallgraph_test.py
new file mode 100644
index 0000000..9561766
--- /dev/null
+++ b/assets/pipegraph/pycallgraph_test.py
@@ -0,0 +1,6 @@
+from pycallgraph2 import PyCallGraph
+from pycallgraph2.output import GraphvizOutput
+from assets.example_test.graph_example import main
+
+with PyCallGraph(output=GraphvizOutput()):
+    main()