fix: filter compr on edge ref

Jaseci-Labs · Jan 27, 2024 · 6ddcc45 · 6ddcc45
1 parent 0bb5adf
commit 6ddcc45
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 17 deletions.
diff --git a/jaclang/core/construct.py b/jaclang/core/construct.py
@@ -41,30 +41,26 @@ def connect_node(self, nd: NodeArchitype, edg: EdgeArchitype) -> NodeArchitype:
         return self.obj
 
     def edges_to_nodes(
-        self, dir: EdgeDir, filter_type: Optional[type],filter_func: Optional[Callable]
+        self, dir: EdgeDir, filter_type: Optional[type], filter_func: Optional[Callable]
     ) -> list[NodeArchitype]:
         """Get set of nodes connected to this node."""
-        print("here", filter_func)
         filter_func = filter_func if filter_func else lambda x: x
         ret_nodes: list[NodeArchitype] = []
-        if dir in [EdgeDir.OUT]: 
-
-            print(self. edges[EdgeDir.OUT])
-            print(filter_func(self. edges[EdgeDir.OUT]))
+        if dir in [EdgeDir.OUT]:
             for i in filter_func(
                 [
                     x
-                    for x in self. edges[EdgeDir.OUT] 
+                    for x in self.edges[EdgeDir.OUT]
                     if x._jac_.target
                     and (not filter_type or isinstance(x, filter_type))
                 ]
             ):
-                ret_nodes. append (i._jac_.target)
+                ret_nodes.append(i._jac_.target)
         elif dir in [EdgeDir.IN]:
             edge = []
-            for i in self.edges[EdgeDir. IN]:
+            for i in self.edges[EdgeDir.IN]:
                 if i._jac_.source and (not filter_type or isinstance(i, filter_type)):
-                    edge.append (i)
+                    edge.append(i)
             new_edge = filter_func(edge)
             for i in new_edge:
                 ret_nodes.append(i._jac_.source)

diff --git a/jaclang/plugin/default.py b/jaclang/plugin/default.py
@@ -129,12 +129,11 @@ def edge_ref(
         node_obj: NodeArchitype,
         dir: EdgeDir,
         filter_type: Optional[type],
-        filter_func: Optional[Callable] = None
+        filter_func: Optional[Callable],
     ) -> list[NodeArchitype]:
         """Jac's apply_dir stmt feature."""
-        print("b",filter_func)
         if isinstance(node_obj, NodeArchitype):
-                return node_obj._jac_.edges_to_nodes(dir, filter_type, filter_func)
+            return node_obj._jac_.edges_to_nodes(dir, filter_type, filter_func)
         else:
             raise TypeError("Invalid node object")
 

diff --git a/jaclang/plugin/feature.py b/jaclang/plugin/feature.py
@@ -86,12 +86,11 @@ def edge_ref(
         node_obj: NodeArchitype,
         dir: EdgeDir,
         filter_type: Optional[type],
-        filter_func: Optional[Callable] = None
+        filter_func: Optional[Callable],
     ) -> list[NodeArchitype]:
         """Jac's apply_dir stmt feature."""
-        print("a", filter_func)
         return JacFeature.pm.hook.edge_ref(
-            node_obj=node_obj, dir=dir, filter_type=filter_type,filter_func=filter_func
+            node_obj=node_obj, dir=dir, filter_type=filter_type, filter_func=filter_func
         )
 
     @staticmethod

diff --git a/jaclang/plugin/spec.py b/jaclang/plugin/spec.py
@@ -110,7 +110,7 @@ def edge_ref(
         node_obj: NodeArchitype,
         dir: EdgeDir,
         filter_type: Optional[type],
-        filter_func: Optional[Callable] = None
+        filter_func: Optional[Callable],
     ) -> list[NodeArchitype]:
         """Jac's apply_dir stmt feature."""
         raise NotImplementedError