Mypy

Author: stellasia

src/neo4j_genai/core/graph.py

@@ -71,29 +71,7 @@ def previous_edges(self, node: Node) -> list[Edge]:
     def __contains__(self, node: Node) -> bool:
         return node.name in self._nodes
 
-    def _check_stack(self, v, visited, stack):
-        visited[v.name] = True
-        stack[v.name] = True
-        for linked_vertices in self.next_edges(v):
-            next_node = linked_vertices.end
-            if not visited[next_node.name]:
-                if self._check_stack(next_node, visited, stack):
-                    return True
-            elif stack[next_node.name]:
-                return True
-        stack[v.name] = False
-        return False
-
-    def __is_cyclic(self) -> bool:
-        visited = {node: False for node in self._nodes}
-        stack = {node: False for node in self._nodes}
-        for node_name, node in self._nodes.items():
-            if not visited[node_name]:
-                if self._check_stack(node, visited, stack):
-                    return True
-        return False
-
-    def dfs(self, visited, node):
+    def dfs(self, visited: set[Node], node: Node) -> bool:
         if node in visited:
             return True
         else:
@@ -103,7 +81,7 @@ def dfs(self, visited, node):
                     return True
             return False
 
-    def is_cyclic(self):
+    def is_cyclic(self) -> bool:
         for node in self._nodes.values():
             if self.dfs(set(), node):
                 return True

src/neo4j_genai/core/pipeline.py

@@ -212,7 +212,7 @@ def get_component_inputs(
                 component_inputs[input_def] = value
         return component_inputs
 
-    def reinitialize(self):
+    def reinitialize(self) -> None:
         self._store.empty()
         self._final_results.empty()
         for task in self._nodes.values():

tests/unit/test_core_graph.py

@@ -2,27 +2,27 @@
 from neo4j_genai.core.graph import Edge, Graph, Node
 
 
-def test_node_alone():
+def test_node_alone() -> None:
     n = Node(name="node", data={})
     assert n.is_root() is True
     assert n.is_leaf() is True
 
 
-def test_node_not_root():
+def test_node_not_root() -> None:
     n = Node(name="node", data={})
     n.parents = [Node(name="child", data={})]
     assert n.is_root() is False
     assert n.is_leaf() is True
 
 
-def test_node_not_leaf():
+def test_node_not_leaf() -> None:
     n = Node(name="node", data={})
     n.children = [Node(name="child", data={})]
     assert n.is_root() is True
     assert n.is_leaf() is False
 
 
-def test_graph_add_nodes():
+def test_graph_add_nodes() -> None:
     g = Graph()
     n1 = Node("n1", {})
     n2 = Node("n2", {})
@@ -36,8 +36,8 @@ def test_graph_add_nodes():
     assert n2.parents == [n1]
 
 
-@pytest.fixture
-def graph():
+@pytest.fixture(scope="function")
+def graph() -> Graph:
     g = Graph()
     n1 = Node("n1", {})
     n2 = Node("n2", {})
@@ -47,13 +47,13 @@ def graph():
     return g
 
 
-def test_graph_roots(graph):
+def test_graph_roots(graph: Graph) -> None:
     roots = graph.roots()
     assert len(roots) == 1
     assert roots[0].name == "n1"
 
 
-def test_graph_next_edge(graph):
+def test_graph_next_edge(graph: Graph) -> None:
     start = graph._nodes["n1"]
     next_edges = graph.next_edges(start)
     assert len(next_edges) == 1
@@ -63,7 +63,7 @@ def test_graph_next_edge(graph):
     assert next_edge.end.name == "n2"
 
 
-def test_graph_prev_edge(graph):
+def test_graph_prev_edge(graph: Graph) -> None:
     start = graph._nodes["n2"]
     next_edges = graph.previous_edges(start)
     assert len(next_edges) == 1
@@ -73,6 +73,6 @@ def test_graph_prev_edge(graph):
     assert next_edge.end.name == "n2"
 
 
-def test_graph_contains(graph):
+def test_graph_contains(graph: Graph) -> None:
     start = graph._nodes["n2"]
     assert start in graph

tests/unit/test_store.py renamed to tests/unit/test_core_store.py

@@ -1,8 +1,8 @@
+import pytest
 from neo4j_genai.core.stores import InMemoryStore
-from sympy.testing import pytest
 
 
-def test_memory_store():
+def test_memory_store() -> None:
     store = InMemoryStore()
     store.add("key", "value")
     assert store.get("key") == "value"