GA-163 | test_digraph (#40)

aMahanna · web-flow · commit 75a3660e90ef · 2024-08-28T10:18:05.000-04:00
* GA-163 | initial commit will fail * unlock adbnx * fix: `incoming_graph_data` * fix: incoming_graph_data * fix: off-by-one IDs * checkpoint * checkpoint: `BaseGraphTester` is passing * checkpoint: BaseGraphAttrTester * cleanup: `aql_fetch_data`, `aql_fetch_data_edge` * use pytest skip for failing tests * checkpoint: optimize `__iter__` * checkpoint: run `test_graph` * add comment * checkpoint * attempt: slleep * fix: lint * cleanup: getitem * cleanup: copy * attempt: shorten sleep * fix: `__set_adj_elements` * fix: mypy * attempt: decrease sleep * GA-163 | `test_digraph` * checkpoint lots of failures... * fix: set `self.Graph` * add type ignore * fix: graph name * fix: graph name * adjust assertions to exclude _rev, set `use_experimental_views` * Revert "adjust assertions to exclude _rev, set `use_experimental_views`" This reverts commit b805419. * fix: `_rev`, `use_experimental_views` * set `use_experimental_views` * fix: lint * new: `nbunch_iter` override * set experimental views to false * set experimental views to false * cleanup * fix: `function.py` * cleanup: `graph`, `digraph` * fix: `test_data_input` * attempt: wait for CircleCI * fix: nx graph * remove sleep * new: `override` suffix * enable more tests * fix: lint
diff --git a/nx_arangodb/classes/dict/adj.py b/nx_arangodb/classes/dict/adj.py
@@ -1557,7 +1557,7 @@ def keys(self) -> Any:
 
     @logger_debug
     def clear(self) -> None:
-        """g._node.clear()"""
+        """g._adj.clear()"""
         self.data.clear()
         self.FETCHED_ALL_DATA = False
         self.FETCHED_ALL_IDS = False
diff --git a/nx_arangodb/classes/digraph.py b/nx_arangodb/classes/digraph.py
@@ -5,9 +5,11 @@
 
 import nx_arangodb as nxadb
 from nx_arangodb.classes.graph import Graph
+from nx_arangodb.logger import logger
 
 from .dict.adj import AdjListOuterDict
 from .enum import TraversalDirection
+from .function import get_node_id
 
 networkx_api = nxadb.utils.decorators.networkx_class(nx.DiGraph)  # type: ignore
 
@@ -59,12 +61,9 @@ def __init__(
         )
 
         if self.graph_exists_in_db:
-            assert isinstance(self._succ, AdjListOuterDict)
-            assert isinstance(self._pred, AdjListOuterDict)
-            self._succ.mirror = self._pred
-            self._pred.mirror = self._succ
-            self._succ.traversal_direction = TraversalDirection.OUTBOUND
-            self._pred.traversal_direction = TraversalDirection.INBOUND
+            self.clear_edges = self.clear_edges_override
+            self.add_node = self.add_node_override
+            self.remove_node = self.remove_node_override
 
     #######################
     # nx.DiGraph Overides #
@@ -80,7 +79,14 @@ def __init__(
     # def out_edges(self):
     # pass
 
-    def add_node(self, node_for_adding, **attr):
+    def clear_edges_override(self):
+        logger.info("Note that clearing edges ony erases the edges in the local cache")
+        for predecessor_dict in self._pred.data.values():
+            predecessor_dict.clear()
+
+        super().clear_edges()
+
+    def add_node_override(self, node_for_adding, **attr):
         if node_for_adding not in self._succ:
             if node_for_adding is None:
                 raise ValueError("None cannot be a node")
@@ -111,7 +117,10 @@ def add_node(self, node_for_adding, **attr):
 
         nx._clear_cache(self)
 
-    def remove_node(self, n):
+    def remove_node_override(self, n):
+        if isinstance(n, (str, int)):
+            n = get_node_id(str(n), self.default_node_type)
+
         try:
 
             ######################
@@ -138,6 +147,22 @@ def remove_node(self, n):
             del self._pred[u][n]  # remove all edges n-u in digraph
         del self._succ[n]  # remove node from succ
         for u in nbrs_pred:
+            ######################
+            # NOTE: Monkey patch #
+            ######################
+
+            # Old: Nothing
+
+            # New:
+            if u == n:
+                continue  # skip self loops
+
+            # Reason: We need to skip self loops, as they are
+            # already taken care of in the previous step. This
+            # avoids getting a KeyError on the next line.
+
+            ###########################
+
             del self._succ[u][n]  # remove all edges n-u in digraph
         del self._pred[n]  # remove node from pred
         nx._clear_cache(self)
diff --git a/nx_arangodb/classes/graph.py b/nx_arangodb/classes/graph.py
@@ -27,6 +27,8 @@
     node_attr_dict_factory,
     node_dict_factory,
 )
+from .dict.adj import AdjListOuterDict
+from .enum import TraversalDirection
 from .function import get_node_id
 from .reportviews import CustomEdgeView, CustomNodeView
 
@@ -96,6 +98,7 @@ def __init__(
         #         m = "Must set **graph_name** if passing **incoming_graph_data**"
         #         raise ValueError(m)
 
+        loaded_incoming_graph_data = False
         if self._graph_exists_in_db:
             if incoming_graph_data is not None:
                 m = "Cannot pass both **incoming_graph_data** and **graph_name** yet if the already graph exists"  # noqa: E501
@@ -170,29 +173,44 @@ def edge_type_func(u: str, v: str) -> str:
                     use_async=write_async,
                 )
 
+                loaded_incoming_graph_data = True
+
             else:
                 self.adb_graph = self.db.create_graph(
                     self.__name,
                     edge_definitions=edge_definitions,
                 )
 
-                # Let the parent class handle the incoming graph data
-                # if it is not a networkx.Graph object
-                kwargs["incoming_graph_data"] = incoming_graph_data
-
             self._set_factory_methods()
             self._set_arangodb_backend_config()
             logger.info(f"Graph '{name}' created.")
             self._graph_exists_in_db = True
 
-        else:
-            kwargs["incoming_graph_data"] = incoming_graph_data
-
-        if name is not None:
-            kwargs["name"] = name
+        if self.__name is not None:
+            kwargs["name"] = self.__name
 
         super().__init__(*args, **kwargs)
 
+        if self.is_directed() and self.graph_exists_in_db:
+            assert isinstance(self._succ, AdjListOuterDict)
+            assert isinstance(self._pred, AdjListOuterDict)
+            self._succ.mirror = self._pred
+            self._pred.mirror = self._succ
+            self._succ.traversal_direction = TraversalDirection.OUTBOUND
+            self._pred.traversal_direction = TraversalDirection.INBOUND
+
+        if incoming_graph_data is not None and not loaded_incoming_graph_data:
+            nx.convert.to_networkx_graph(incoming_graph_data, create_using=self)
+
+        if self.graph_exists_in_db:
+            self.copy = self.copy_override
+            self.subgraph = self.subgraph_override
+            self.clear = self.clear_override
+            self.clear_edges = self.clear_edges_override
+            self.add_node = self.add_node_override
+            self.number_of_edges = self.number_of_edges_override
+            self.nbunch_iter = self.nbunch_iter_override
+
     #######################
     # Init helper methods #
     #######################
@@ -345,6 +363,9 @@ def _set_graph_name(self, graph_name: str | None = None) -> None:
     # ArangoDB Methods #
     ####################
 
+    def clear_nxcg_cache(self):
+        self.nxcg_graph = None
+
     def aql(self, query: str, bind_vars: dict[str, Any] = {}, **kwargs: Any) -> Cursor:
         return nxadb.classes.function.aql(self.db, query, bind_vars, **kwargs)
 
@@ -355,7 +376,7 @@ def aql(self, query: str, bind_vars: dict[str, Any] = {}, **kwargs: Any) -> Curs
     # NOTE: OUT OF SERVICE
     # def chat(self, prompt: str) -> str:
     #     if self.__qa_chain is None:
-    #         if not self.__graph_exists_in_db:
+    #         if not self.graph_exists_in_db:
     #             return "Could not initialize QA chain: Graph does not exist"
 
     #         # try:
@@ -381,30 +402,6 @@ def aql(self, query: str, bind_vars: dict[str, Any] = {}, **kwargs: Any) -> Curs
     # nx.Graph Overides #
     #####################
 
-    def copy(self, *args, **kwargs):
-        logger.warning("Note that copying a graph loses the connection to the database")
-        G = super().copy(*args, **kwargs)
-        G.node_dict_factory = nx.Graph.node_dict_factory
-        G.node_attr_dict_factory = nx.Graph.node_attr_dict_factory
-        G.edge_attr_dict_factory = nx.Graph.edge_attr_dict_factory
-        G.adjlist_inner_dict_factory = nx.Graph.adjlist_inner_dict_factory
-        G.adjlist_outer_dict_factory = nx.Graph.adjlist_outer_dict_factory
-        return G
-
-    def subgraph(self, nbunch):
-        raise NotImplementedError("Subgraphing is not yet implemented")
-
-    def clear(self):
-        logger.info("Note that clearing only erases the local cache")
-        super().clear()
-
-    def clear_edges(self):
-        logger.info("Note that clearing edges ony erases the edges in the local cache")
-        super().clear_edges()
-
-    def clear_nxcg_cache(self):
-        self.nxcg_graph = None
-
     @cached_property
     def nodes(self):
         if self.__use_experimental_views and self.graph_exists_in_db:
@@ -437,7 +434,30 @@ def edges(self):
 
         return super().edges
 
-    def add_node(self, node_for_adding, **attr):
+    def copy_override(self, *args, **kwargs):
+        logger.warning("Note that copying a graph loses the connection to the database")
+        G = super().copy(*args, **kwargs)
+        G.node_dict_factory = nx.Graph.node_dict_factory
+        G.node_attr_dict_factory = nx.Graph.node_attr_dict_factory
+        G.edge_attr_dict_factory = nx.Graph.edge_attr_dict_factory
+        G.adjlist_inner_dict_factory = nx.Graph.adjlist_inner_dict_factory
+        G.adjlist_outer_dict_factory = nx.Graph.adjlist_outer_dict_factory
+        return G
+
+    def subgraph_override(self, nbunch):
+        raise NotImplementedError("Subgraphing is not yet implemented")
+
+    def clear_override(self):
+        logger.info("Note that clearing only erases the local cache")
+        super().clear()
+
+    def clear_edges_override(self):
+        logger.info("Note that clearing edges ony erases the edges in the local cache")
+        for nbr_dict in self._adj.data.values():
+            nbr_dict.clear()
+        nx._clear_cache(self)
+
+    def add_node_override(self, node_for_adding, **attr):
         if node_for_adding not in self._node:
             if node_for_adding is None:
                 raise ValueError("None cannot be a node")
@@ -467,10 +487,7 @@ def add_node(self, node_for_adding, **attr):
 
         nx._clear_cache(self)
 
-    def number_of_edges(self, u=None, v=None):
-        if not self.graph_exists_in_db:
-            return super().number_of_edges(u, v)
-
+    def number_of_edges_override(self, u=None, v=None):
         if u is not None:
             return super().number_of_edges(u, v)
 
@@ -494,10 +511,7 @@ def number_of_edges(self, u=None, v=None):
         # It is more efficient to count the number of edges in the edge collections
         # compared to relying on the DegreeView.
 
-    def nbunch_iter(self, nbunch=None):
-        if not self._graph_exists_in_db:
-            return super().nbunch_iter(nbunch)
-
+    def nbunch_iter_override(self, nbunch=None):
         if nbunch is None:
             bunch = iter(self._adj)
         elif nbunch in self:
@@ -508,12 +522,13 @@ def nbunch_iter(self, nbunch=None):
             # Old: Nothing
 
             # New:
-            if isinstance(nbunch, int):
+            if isinstance(nbunch, (str, int)):
                 nbunch = get_node_id(str(nbunch), self.default_node_type)
 
             # Reason:
             # ArangoDB only uses strings as node IDs. Therefore, we need to convert
-            # the integer node ID to a string before using it in an iterator.
+            # the non-prefixed node ID to an ArangoDB ID before
+            # using it in an iterator.
 
             bunch = iter([nbunch])
         else:
@@ -528,13 +543,15 @@ def bunch_iter(nlist, adj):
                         # Old: Nothing
 
                         # New:
-                        if isinstance(n, int):
+                        if isinstance(n, (str, int)):
                             n = get_node_id(str(n), self.default_node_type)
 
                         # Reason:
                         # ArangoDB only uses strings as node IDs. Therefore,
-                        # we need to convert the integer node ID to a
-                        # string before using it in an iterator.
+                        # we need to convert non-prefixed node IDs to an
+                        # ArangoDB ID before using it in an iterator.
+
+                        ######################
 
                         if n in adj:
                             yield n
diff --git a/nx_arangodb/classes/multigraph.py b/nx_arangodb/classes/multigraph.py
@@ -58,6 +58,9 @@ def __init__(
             **kwargs,
         )
 
+        if self.graph_exists_in_db:
+            self.add_edge = self.add_edge_override
+
     #######################
     # Init helper methods #
     #######################
@@ -76,10 +79,7 @@ def _set_factory_methods(self) -> None:
     # nx.MultiGraph Overides #
     ##########################
 
-    def add_edge(self, u_for_edge, v_for_edge, key=None, **attr):
-        if not self.graph_exists_in_db:
-            return super().add_edge(u_for_edge, v_for_edge, key=key, **attr)
-
+    def add_edge_override(self, u_for_edge, v_for_edge, key=None, **attr):
         if key is not None:
             m = "ArangoDB MultiGraph does not support custom edge keys yet."
             logger.warning(m)
diff --git a/tests/test.py b/tests/test.py
@@ -18,6 +18,9 @@
 from .conftest import Capturing, create_grid_graph, create_line_graph, db, run_gpu_tests
 
 G_NX = nx.karate_club_graph()
+G_NX_digraph = nx.DiGraph(G_NX)
+G_NX_multigraph = nx.MultiGraph(G_NX)
+G_NX_multidigraph = nx.MultiDiGraph(G_NX)
 
 
 def assert_remote_dict(G: nxadb.Graph) -> None:
@@ -1834,3 +1837,58 @@ def test_incoming_graph_data_not_nx_graph(
     )
     assert has_club == ("club" in G.nodes["0"])
     assert has_weight == ("weight" in G.adj["0"]["1"])
+
+
+@pytest.mark.parametrize(
+    "data_type, incoming_graph_data, has_club, has_weight",
+    [
+        ("dict of dicts", G_NX_digraph._adj, False, True),
+        (
+            "dict of lists",
+            {k: list(v) for k, v in G_NX_digraph._adj.items()},
+            False,
+            False,
+        ),
+        ("container of edges", list(G_NX_digraph.edges), False, False),
+        ("iterator of edges", iter(G_NX_digraph.edges), False, False),
+        ("generator of edges", (e for e in G_NX_digraph.edges), False, False),
+        ("2D numpy array", nx.to_numpy_array(G_NX_digraph), False, True),
+        (
+            "scipy sparse array",
+            nx.to_scipy_sparse_array(G_NX_digraph),
+            False,
+            True,
+        ),
+        ("Pandas EdgeList", nx.to_pandas_edgelist(G_NX_digraph), False, True),
+        ("Pandas Adjacency", nx.to_pandas_adjacency(G_NX_digraph), False, True),
+    ],
+)
+def test_incoming_graph_data_not_nx_graph_digraph(
+    data_type: str, incoming_graph_data: Any, has_club: bool, has_weight: bool
+) -> None:
+    # See nx.convert.to_networkx_graph for the official supported types
+    name = "KarateGraph"
+    db.delete_graph(name, drop_collections=True, ignore_missing=True)
+
+    G = nxadb.DiGraph(incoming_graph_data=incoming_graph_data, name=name)
+
+    assert (
+        len(G.adj)
+        == len(G_NX_digraph.adj)
+        == db.collection(G.default_node_type).count()
+    )
+    assert (
+        len(G.nodes)
+        == len(G_NX_digraph.nodes)
+        == db.collection(G.default_node_type).count()
+        == G.number_of_nodes()
+    )
+    edge_col = G.edge_type_func(G.default_node_type, G.default_node_type)
+    assert (
+        len(G.edges)
+        == len(G_NX_digraph.edges)
+        == db.collection(edge_col).count()
+        == G.number_of_edges()
+    )
+    assert has_club == ("club" in G.nodes["0"])
+    assert has_weight == ("weight" in G.adj["0"]["1"])
diff --git a/tests/test_digraph.py b/tests/test_digraph.py
diff --git a/tests/test_graph.py b/tests/test_graph.py
