replace tqdm with rich

arangoml · Jul 22, 2022 · 8dfbb75 · 8dfbb75
1 parent 5b56112
commit 8dfbb75
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 26 deletions.
diff --git a/adbnx_adapter/adapter.py b/adbnx_adapter/adapter.py
@@ -10,7 +10,8 @@
 from arango.result import Result
 from networkx.classes.graph import Graph as NXGraph
 from networkx.classes.multidigraph import MultiDiGraph as NXMultiDiGraph
-from tqdm import tqdm
+from rich.progress import track
+from rich.status import Status
 
 from .abc import Abstract_ADBNX_Adapter
 from .controller import ADBNX_Controller
@@ -121,13 +122,15 @@ def arangodb_to_networkx(
         for v_col, atribs in metagraph["vertexCollections"].items():
             logger.debug(f"Preparing '{v_col}' vertices")
 
+            nx_nodes.clear()
             cursor = self.__fetch_adb_docs(v_col, atribs, is_keep, query_options)
-            for adb_v in tqdm(
+            for adb_v in track(
                 cursor,
                 total=cursor.count(),
-                desc=v_col,
-                colour="CYAN",
-                disable=logger.level > logging.INFO,
+                description=v_col,
+                complete_style="#079DE8",
+                finished_style="#079DE8",
+                disable=logger.level != logging.INFO,
             ):
                 adb_id: str = adb_v["_id"]
                 self.__cntrl._prepare_arangodb_vertex(adb_v, v_col)
@@ -138,19 +141,19 @@ def arangodb_to_networkx(
 
             logger.debug(f"Inserting {len(nx_nodes)} '{v_col}' vertices")
             nx_graph.add_nodes_from(nx_nodes)
-            nx_nodes.clear()
 
         adb_e: Json
         for e_col, atribs in metagraph["edgeCollections"].items():
             logger.debug(f"Preparing '{e_col}' edges")
 
             cursor = self.__fetch_adb_docs(e_col, atribs, is_keep, query_options)
-            for adb_e in tqdm(
+            for adb_e in track(
                 cursor,
                 total=cursor.count(),
-                desc=e_col,
-                colour="GREEN",
-                disable=logger.level > logging.INFO,
+                description=e_col,
+                complete_style="#FA7D05",
+                finished_style="#FA7D05",
+                disable=logger.level != logging.INFO,
             ):
                 from_node_id: NxId = adb_map[adb_e["_from"]]
                 to_node_id: NxId = adb_map[adb_e["_to"]]
@@ -160,7 +163,6 @@ def arangodb_to_networkx(
 
             logger.debug(f"Inserting {len(nx_edges)} '{e_col}' edges")
             nx_graph.add_edges_from(nx_edges)
-            nx_edges.clear()
 
         logger.info(f"Created NetworkX '{name}' Graph")
         return nx_graph
@@ -278,8 +280,10 @@ def networkx_to_arangodb(
             self.__db.delete_graph(name, ignore_missing=True)
 
         if self.__db.has_graph(name):
+            logger.debug(f"Graph {name} already exists")
             adb_graph = self.__db.graph(name)
         else:
+            logger.debug(f"Creating graph {name}")
             adb_graph = self.__db.create_graph(name, edge_definitions)
 
         adb_v_cols: List[str] = adb_graph.vertex_collections()
@@ -301,14 +305,17 @@ def networkx_to_arangodb(
         nx_node: NxData
         logger.debug("Preparing NetworkX nodes")
         for i, (nx_id, nx_node) in enumerate(
-            tqdm(
+            track(
                 nx_graph.nodes(data=True),
-                desc="Nodes",
-                colour="CYAN",
-                disable=logger.level > logging.INFO,
+                description="Nodes",
+                complete_style="#97C423",
+                finished_style="#97C423",
+                disable=logger.level != logging.INFO,
             ),
             1,
         ):
+            logger.debug(f"N{i}: {nx_id}")
+
             col = (
                 adb_v_cols[0]
                 if has_one_vcol
@@ -343,14 +350,18 @@ def networkx_to_arangodb(
         nx_edge: NxData
         logger.debug("Preparing NetworkX edges")
         for i, (from_node_id, to_node_id, nx_edge) in enumerate(
-            tqdm(
+            track(
                 nx_graph.edges(data=True),
-                desc="Edges",
-                colour="GREEN",
-                disable=logger.level > logging.INFO,
+                description="Edges",
+                complete_style="#5E3108",
+                finished_style="#5E3108",
+                disable=logger.level != logging.INFO,
             ),
             1,
         ):
+            edge_str = f"({from_node_id}, {to_node_id})"
+            logger.debug(f"E{i}: {edge_str}")
+
             from_n = nx_map[from_node_id]
             to_n = nx_map[to_node_id]
 
@@ -363,7 +374,6 @@ def networkx_to_arangodb(
             )
 
             if col not in adb_e_cols:
-                edge_str = f"({from_node_id}, {to_node_id})"
                 msg = f"{edge_str} identified as '{col}', which is not in {adb_e_cols}"
                 raise ValueError(msg)
 
@@ -425,17 +435,21 @@ def __fetch_adb_docs(
         )
 
     def __insert_adb_docs(
-        self, adb_documents: DefaultDict[str, List[Json]], import_options: Any
+        self, adb_documents: DefaultDict[str, List[Json]], kwargs: Any
     ) -> None:
         """Insert ArangoDB documents into their ArangoDB collection.
 
         :param adb_documents: To-be-inserted ArangoDB documents
         :type adb_documents: DefaultDict[str, List[Json]]
-        :param import_options: Keyword arguments to specify additional
+        :param kwargs: Keyword arguments to specify additional
             parameters for ArangoDB document insertion. Full parameter list:
             https://docs.python-arango.com/en/main/specs.html#arango.collection.Collection.import_bulk
         """
         for col, doc_list in adb_documents.items():
-            logger.debug(f"Inserting {len(doc_list)} documents into '{col}'")
-            result = self.__db.collection(col).import_bulk(doc_list, **import_options)
-            logger.debug(result)
+            with Status(
+                f"POST /_api/import '{col}' ({len(doc_list)})",
+                spinner="aesthetic",
+                spinner_style="cyan",
+            ):
+                result = self.__db.collection(col).import_bulk(doc_list, **kwargs)
+                logger.debug(result)
diff --git a/setup.py b/setup.py
@@ -20,7 +20,7 @@
         "requests>=2.27.1",
         "python-arango>=7.4.1",
         "networkx>=2.5.1",
-        "tqdm>=4.64.0",
+        "rich>=12.5.1",
         "setuptools>=45",
     ],
     extras_require={