biocypher 0.5.41__py3-none-any.whl → 0.5.43__py3-none-any.whl

biocypher/output/write/_writer.py

@@ -0,0 +1,200 @@
+from abc import ABC, abstractmethod
+from typing import Union, Optional
+from collections.abc import Iterable
+import os
+
+from biocypher._create import BioCypherEdge, BioCypherNode, BioCypherRelAsNode
+from biocypher._logger import logger
+from biocypher._translate import Translator
+from biocypher._deduplicate import Deduplicator
+
+__all__ = ["_Writer"]
+
+
+class _Writer(ABC):
+    """Abstract class for writing node and edge representations to disk.
+    Specifics of the different writers (e.g. neo4j, postgresql, csv, etc.)
+    are implemented in the child classes. Any concrete writer needs to
+    implement at least:
+    - _write_node_data
+    - _write_edge_data
+    - _construct_import_call
+    - _get_import_script_name
+
+    Args:
+        translator (Translator): Instance of :py:class:`Translator` to enable translation of
+            nodes and manipulation of properties.
+        deduplicator (Deduplicator): Instance of :py:class:`Deduplicator` to enable deduplication
+            of nodes and edges.
+        output_directory (str, optional): Path for exporting CSV files. Defaults to None.
+        strict_mode (bool, optional): Whether to enforce source, version, and license properties. Defaults to False.
+    strict_mode (bool, optional): Whether to enforce source, version, and license properties. Defaults to False.
+
+    Raises:
+        NotImplementedError: Writer implementation must override '_write_node_data'
+        NotImplementedError: Writer implementation must override '_write_edge_data'
+        NotImplementedError: Writer implementation must override '_construct_import_call'
+        NotImplementedError: Writer implementation must override '_get_import_script_name'
+    """
+
+    def __init__(
+        self,
+        translator: Translator,
+        deduplicator: Deduplicator,
+        output_directory: Optional[str] = None,
+        strict_mode: bool = False,
+        *args,
+        **kwargs,
+    ):
+        """Abstract class for writing node and edge representations to disk.
+
+        Args:
+            translator (Translator): Instance of :py:class:`Translator` to enable translation of
+                nodes and manipulation of properties.
+            deduplicator (Deduplicator): Instance of :py:class:`Deduplicator` to enable deduplication
+                of nodes and edges.
+            output_directory (str, optional): Path for exporting CSV files. Defaults to None.
+            strict_mode (bool, optional): Whether to enforce source, version, and license properties. Defaults to False.
+        strict_mode (bool, optional): Whether to enforce source, version, and license properties. Defaults to False.
+        """
+        self.translator = translator
+        self.deduplicator = deduplicator
+        self.strict_mode = strict_mode
+        self.output_directory = output_directory
+
+        if os.path.exists(self.output_directory):
+            if kwargs.get("write_to_file", True):
+                logger.warning(
+                    f"Output directory `{self.output_directory}` already exists. "
+                    "If this is not planned, file consistency may be compromised."
+                )
+        else:
+            logger.info(f"Creating output directory `{self.output_directory}`.")
+            os.makedirs(self.output_directory)
+
+    @abstractmethod
+    def _write_node_data(
+        self,
+        nodes: Iterable[
+            Union[BioCypherNode, BioCypherEdge, BioCypherRelAsNode]
+        ],
+    ) -> bool:
+        """Implement how to output.write nodes to disk.
+
+        Args:
+            nodes (Iterable): An iterable of BioCypherNode / BioCypherEdge / BioCypherRelAsNode objects.
+
+        Returns:
+            bool: The return value. True for success, False otherwise.
+        """
+        raise NotImplementedError(
+            "Writer implementation must override 'write_nodes'"
+        )
+
+    @abstractmethod
+    def _write_edge_data(
+        self,
+        edges: Iterable[
+            Union[BioCypherNode, BioCypherEdge, BioCypherRelAsNode]
+        ],
+    ) -> bool:
+        """Implement how to output.write edges to disk.
+
+        Args:
+            edges (Iterable): An iterable of BioCypherNode / BioCypherEdge / BioCypherRelAsNode objects.
+
+        Returns:
+            bool: The return value. True for success, False otherwise.
+        """
+        raise NotImplementedError(
+            "Writer implementation must override 'write_edges'"
+        )
+
+    @abstractmethod
+    def _construct_import_call(self) -> str:
+        """
+        Function to construct the import call detailing folder and
+        individual node and edge headers and data files, as well as
+        delimiters and database name. Built after all data has been
+        processed to ensure that nodes are called before any edges.
+
+        Returns:
+            str: command for importing the output files into a DBMS.
+        """
+        raise NotImplementedError(
+            "Writer implementation must override '_construct_import_call'"
+        )
+
+    @abstractmethod
+    def _get_import_script_name(self) -> str:
+        """Returns the name of the import script.
+
+        Returns:
+            str: The name of the import script (ending in .sh)
+        """
+        raise NotImplementedError(
+            "Writer implementation must override '_get_import_script_name'"
+        )
+
+    def write_nodes(
+        self, nodes, batch_size: int = int(1e6), force: bool = False
+    ):
+        """Wrapper for writing nodes.
+
+        Args:
+            nodes (BioCypherNode): a list or generator of nodes in
+                :py:class:`BioCypherNode` format
+            batch_size (int): The batch size for writing nodes.
+            force (bool): Whether to force writing nodes even if their type is
+                not present in the schema.
+
+        Returns:
+            bool: The return value. True for success, False otherwise.
+        """
+        passed = self._write_node_data(nodes)
+        if not passed:
+            logger.error("Error while writing node data.")
+            return False
+        return True
+
+    def write_edges(
+        self, edges, batch_size: int = int(1e6), force: bool = False
+    ):
+        """Wrapper for writing edges.
+
+        Args:
+            nodes (BioCypherNode): a list or generator of nodes in
+                :py:class:`BioCypherNode` format
+            batch_size (int): The batch size for writing nodes.
+            force (bool): Whether to force writing nodes even if their type is
+                not present in the schema.
+
+        Returns:
+            bool: The return value. True for success, False otherwise.
+        """
+        passed = self._write_edge_data(edges)
+        if not passed:
+            logger.error("Error while writing edge data.")
+            return False
+        return True
+
+    def write_import_call(self):
+        """
+        Function to output.write the import call detailing folder and
+        individual node and edge headers and data files, as well as
+        delimiters and database name, to the export folder as txt.
+
+        Returns:
+            str: The path of the file holding the import call.
+        """
+        file_path = os.path.join(
+            self.output_directory, self._get_import_script_name()
+        )
+        logger.info(
+            f"Writing {self.__class__.__name__} import call to `{file_path}`."
+        )
+
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(self._construct_import_call())
+
+        return file_path

biocypher/{write → output/write}/graph/_arangodb.py

@@ -1,7 +1,7 @@
 import os
 
 from biocypher._logger import logger
-from biocypher.write.graph._neo4j import _Neo4jBatchWriter
+from biocypher.output.write.graph._neo4j import _Neo4jBatchWriter
 
 
 class _ArangoDBBatchWriter(_Neo4jBatchWriter):

biocypher/{write → output/write}/graph/_neo4j.py

@@ -1,9 +1,7 @@
 import os
-import re
-import subprocess
 
 from biocypher._logger import logger
-from biocypher.write._batch_writer import parse_label, _BatchWriter
+from biocypher.output.write._batch_writer import parse_label, _BatchWriter
 
 
 class _Neo4jBatchWriter(_BatchWriter):
@@ -49,7 +47,7 @@ class _Neo4jBatchWriter(_BatchWriter):
 
     def _write_array_string(self, string_list):
         """
-        Abstract method to write the string representation of an array into a .csv file
+        Abstract method to output.write the string representation of an array into a .csv file
         as required by the neo4j admin-import.
 
         Args:

biocypher/output/write/graph/_networkx.py

@@ -0,0 +1,76 @@
+import pickle
+
+import networkx as nx
+
+from biocypher._logger import logger
+from biocypher.output.write._writer import _Writer
+from biocypher.output.write.relational._csv import _PandasCSVWriter
+
+
+class _NetworkXWriter(_Writer):
+    """
+    Class for writing node and edges to a networkx DiGraph.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.csv_writer = _PandasCSVWriter(*args, write_to_file=False, **kwargs)
+        self.G = nx.DiGraph()
+
+    def _construct_import_call(self) -> str:
+        """Function to construct the Python code to load all node and edge csv files again into Pandas dfs.
+
+        Returns:
+            str: Python code to load the csv files into Pandas dfs.
+        """
+        logger.info(
+            f"Writing networkx {self.G} to pickle file networkx_graph.pkl."
+        )
+        with open(f"{self.output_directory}/networkx_graph.pkl", "wb") as f:
+            pickle.dump(self.G, f)
+
+        import_call = "import pickle\n"
+        import_call += "with open('./networkx_graph.pkl', 'rb') as f:\n\tG_loaded = pickle.load(f)"
+        return import_call
+
+    def _get_import_script_name(self) -> str:
+        """Function to return the name of the import script."""
+        return "import_networkx.py"
+
+    def _write_node_data(self, nodes) -> bool:
+        passed = self.csv_writer._write_entities_to_file(nodes)
+        self.add_to_networkx()
+        return passed
+
+    def _write_edge_data(self, edges) -> bool:
+        passed = self.csv_writer._write_entities_to_file(edges)
+        self.add_to_networkx()
+        return passed
+
+    def add_to_networkx(self) -> bool:
+        all_dfs = self.csv_writer.stored_dfs
+        node_dfs = [
+            df
+            for df in all_dfs.values()
+            if df.columns.str.contains("node_id").any()
+        ]
+        edge_dfs = [
+            df
+            for df in all_dfs.values()
+            if df.columns.str.contains("source_id").any()
+            and df.columns.str.contains("target_id").any()
+        ]
+        for df in node_dfs:
+            nodes = df.set_index("node_id").to_dict(orient="index")
+            self.G.add_nodes_from(nodes.items())
+        for df in edge_dfs:
+            edges = df.set_index(["source_id", "target_id"]).to_dict(
+                orient="index"
+            )
+            self.G.add_edges_from(
+                (
+                    (source, target, attrs)
+                    for (source, target), attrs in edges.items()
+                )
+            )
+        return True