diff --git a/adbnx_adapter/adapter.py b/adbnx_adapter/adapter.py index df5dc1c0..1d466f7e 100644 --- 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 index f9abd53b..d8ffabfa 100644 --- 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={