biocypher 0.6.2__py3-none-any.whl → 0.7.0__py3-none-any.whl

biocypher/output/write/_batch_writer.py

@@ -1,17 +1,17 @@
-from abc import ABC, abstractmethod
-from types import GeneratorType
-from typing import Union, Optional
-from collections import OrderedDict, defaultdict
+import glob
 import os
 import re
-import glob
+
+from abc import ABC, abstractmethod
+from collections import OrderedDict, defaultdict
+from types import GeneratorType
 
 from more_itertools import peekable
 
 from biocypher._create import BioCypherEdge, BioCypherNode, BioCypherRelAsNode
+from biocypher._deduplicate import Deduplicator
 from biocypher._logger import logger
 from biocypher._translate import Translator
-from biocypher._deduplicate import Deduplicator
 from biocypher.output.write._writer import _Writer
 
 
@@ -20,96 +20,88 @@ class _BatchWriter(_Writer, ABC):
 
     @abstractmethod
     def _quote_string(self, value: str) -> str:
-        """
-        Abstract method to quote a string. Escaping is handled by the database-specific writer.
-        """
-
+        """Abstract method to quote a string. Escaping is handled by the database-specific writer."""
         raise NotImplementedError(
-            "Database writer must override '_quote_string'"
+            "Database writer must override '_quote_string'",
         )
 
     @abstractmethod
     def _get_default_import_call_bin_prefix(self):
-        """
-        Abstract method to provide the default string for the import call bin prefix.
+        """Abstract method to provide the default string for the import call bin prefix.
 
-        Returns:
+        Returns
+        -------
             str: The database-specific string for the path to the import call bin prefix
+
         """
-        raise NotImplementedError(
-            "Database writer must override '_get_default_import_call_bin_prefix'"
-        )
+        raise NotImplementedError("Database writer must override '_get_default_import_call_bin_prefix'")
 
     @abstractmethod
     def _write_array_string(self, string_list):
-        """
-        Abstract method to write the string representation of an array into a .csv file.
+        """Abstract method to write the string representation of an array into a .csv file.
         Different databases require different formats of array to optimize import speed.
 
         Args:
+        ----
             string_list (list): list of ontology strings
 
         Returns:
+        -------
             str: The database-specific string representation of an array
+
         """
-        raise NotImplementedError(
-            "Database writer must override '_write_array_string'"
-        )
+        raise NotImplementedError("Database writer must override '_write_array_string'")
 
     @abstractmethod
     def _write_node_headers(self):
-        """
-        Abstract method that takes care of importing properties of a graph entity that is represented
+        """Abstract method that takes care of importing properties of a graph entity that is represented
         as a node as per the definition in the `schema_config.yaml`
 
-        Returns:
+        Returns
+        -------
             bool: The return value. True for success, False otherwise.
+
         """
-        raise NotImplementedError(
-            "Database writer must override '_write_node_headers'"
-        )
+        raise NotImplementedError("Database writer must override '_write_node_headers'")
 
     @abstractmethod
     def _write_edge_headers(self):
-        """
-        Abstract method to write a database import-file for a graph entity that is represented
+        """Abstract method to write a database import-file for a graph entity that is represented
         as an edge as per the definition in the `schema_config.yaml`,
         containing only the header for this type of edge.
 
-        Returns:
+        Returns
+        -------
             bool: The return value. True for success, False otherwise.
+
         """
-        raise NotImplementedError(
-            "Database writer must override '_write_edge_headers'"
-        )
+        raise NotImplementedError("Database writer must override '_write_edge_headers'")
 
     @abstractmethod
     def _construct_import_call(self) -> str:
-        """
-        Function to construct the import call detailing folder and
+        """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:
+        Returns
+        -------
             str: A bash command for csv import.
+
         """
-        raise NotImplementedError(
-            "Database writer must override '_construct_import_call'"
-        )
+        raise NotImplementedError("Database writer must override '_construct_import_call'")
 
     @abstractmethod
     def _get_import_script_name(self) -> str:
-        """
-        Returns the name of the import script.
+        """Returns the name of the import script.
         The name will be chosen based on the used database.
 
-        Returns:
+        Returns
+        -------
             str: The name of the import script (ending in .sh)
+
         """
-        raise NotImplementedError(
-            "Database writer must override '_get_import_script_name'"
-        )
+        raise NotImplementedError("Database writer must override '_get_import_script_name'")
 
     def __init__(
         self,
@@ -118,10 +110,10 @@ class _BatchWriter(_Writer, ABC):
         delimiter: str,
         array_delimiter: str = ",",
         quote: str = '"',
-        output_directory: Optional[str] = None,
+        output_directory: str | None = None,
         db_name: str = "neo4j",
-        import_call_bin_prefix: Optional[str] = None,
-        import_call_file_prefix: Optional[str] = None,
+        import_call_bin_prefix: str | None = None,
+        import_call_file_prefix: str | None = None,
         wipe: bool = True,
         strict_mode: bool = False,
         skip_bad_relationships: bool = False,
@@ -133,9 +125,7 @@ class _BatchWriter(_Writer, ABC):
         rdf_format: str = None,
         rdf_namespaces: dict = {},
     ):
-        """
-
-        Abtract parent class for writing node and edge representations to disk
+        """Abtract parent class for writing node and edge representations to disk
         using the format specified by each database type. The database-specific
         functions are implemented by the respective child-classes. This abstract
         class contains all methods expected by a bach writer instance, some of
@@ -156,6 +146,7 @@ class _BatchWriter(_Writer, ABC):
             - _get_import_script_name
 
         Args:
+        ----
             translator:
                 Instance of :py:class:`Translator` to enable translation of
                 nodes and manipulation of properties.
@@ -217,6 +208,7 @@ class _BatchWriter(_Writer, ABC):
 
             rdf_namespaces:
                 The namespaces for RDF.
+
         """
         super().__init__(
             translator=translator,
@@ -233,17 +225,13 @@ class _BatchWriter(_Writer, ABC):
         self.rdf_namespaces = rdf_namespaces
 
         self.delim, self.escaped_delim = self._process_delimiter(delimiter)
-        self.adelim, self.escaped_adelim = self._process_delimiter(
-            array_delimiter
-        )
+        self.adelim, self.escaped_adelim = self._process_delimiter(array_delimiter)
         self.quote = quote
         self.skip_bad_relationships = skip_bad_relationships
         self.skip_duplicate_nodes = skip_duplicate_nodes
 
         if import_call_bin_prefix is None:
-            self.import_call_bin_prefix = (
-                self._get_default_import_call_bin_prefix()
-            )
+            self.import_call_bin_prefix = self._get_default_import_call_bin_prefix()
         else:
             self.import_call_bin_prefix = import_call_bin_prefix
 
@@ -268,34 +256,27 @@ class _BatchWriter(_Writer, ABC):
 
     @property
     def import_call_file_prefix(self):
-        """
-        Property for output directory path.
-        """
-
+        """Property for output directory path."""
         if self._import_call_file_prefix is None:
             return self.outdir
         else:
             return self._import_call_file_prefix
 
     def _process_delimiter(self, delimiter: str) -> str:
-        """
-        Return escaped characters in case of receiving their string
+        """Return escaped characters in case of receiving their string
         representation (e.g. tab for '\t').
         """
-
         if delimiter == "\\t":
             return "\t", "\\t"
 
         else:
             return delimiter, delimiter
 
-    def write_nodes(
-        self, nodes, batch_size: int = int(1e6), force: bool = False
-    ):
-        """
-        Wrapper for writing nodes and their headers.
+    def write_nodes(self, nodes, batch_size: int = int(1e6), force: bool = False):
+        """Wrapper for writing nodes and their headers.
 
         Args:
+        ----
             nodes (BioCypherNode): a list or generator of nodes in
                 :py:class:`BioCypherNode` format
 
@@ -306,7 +287,9 @@ class _BatchWriter(_Writer, ABC):
 
 
         Returns:
+        -------
             bool: The return value. True for success, False otherwise.
+
         """
         # TODO check represented_as
 
@@ -325,19 +308,21 @@ class _BatchWriter(_Writer, ABC):
 
     def write_edges(
         self,
-        edges: Union[list, GeneratorType],
+        edges: list | GeneratorType,
         batch_size: int = int(1e6),
     ) -> bool:
-        """
-        Wrapper for writing edges and their headers.
+        """Wrapper for writing edges and their headers.
 
         Args:
+        ----
             edges (BioCypherEdge): a list or generator of edges in
                 :py:class:`BioCypherEdge` or :py:class:`BioCypherRelAsNode`
                 format
 
         Returns:
+        -------
             bool: The return value. True for success, False otherwise.
+
         """
         passed = False
         edges = list(edges)  # force evaluation to handle empty generator
@@ -375,7 +360,6 @@ class _BatchWriter(_Writer, ABC):
             logger.debug(
                 "No edges to write, possibly due to no matched Biolink classes.",
             )
-            pass
 
         if not passed:
             logger.error("Error while writing edge data.")
@@ -389,8 +373,7 @@ class _BatchWriter(_Writer, ABC):
         return True
 
     def _write_node_data(self, nodes, batch_size, force: bool = False):
-        """
-        Writes biocypher nodes to CSV conforming to the headers created
+        """Writes biocypher nodes to CSV conforming to the headers created
         with `_write_node_headers()`, and is actually required to be run
         before calling `_write_node_headers()` to set the
         :py:attr:`self.node_property_dict` for passing the node properties
@@ -398,14 +381,16 @@ class _BatchWriter(_Writer, ABC):
         :py:class:`BioCypherNode` class.
 
         Args:
+        ----
             nodes (BioCypherNode): a list or generator of nodes in
                 :py:class:`BioCypherNode` format
 
         Returns:
+        -------
             bool: The return value. True for success, False otherwise.
-        """
 
-        if isinstance(nodes, GeneratorType) or isinstance(nodes, peekable):
+        """
+        if isinstance(nodes, GeneratorType | peekable):
             logger.debug("Writing node CSV from generator.")
 
             bins = defaultdict(list)  # dict to store a list for each
@@ -432,20 +417,15 @@ class _BatchWriter(_Writer, ABC):
                     logger.warning(f"Node {label} has no id; skipping.")
                     continue
 
-                if not label in bins.keys():
+                if label not in bins.keys():
                     # start new list
                     all_labels = None
                     bins[label].append(node)
                     bin_l[label] = 1
 
                     # get properties from config if present
-                    if (
-                        label
-                        in self.translator.ontology.mapping.extended_schema
-                    ):
-                        cprops = self.translator.ontology.mapping.extended_schema.get(
-                            label
-                        ).get(
+                    if label in self.translator.ontology.mapping.extended_schema:
+                        cprops = self.translator.ontology.mapping.extended_schema.get(label).get(
                             "properties",
                         )
                     else:
@@ -483,18 +463,13 @@ class _BatchWriter(_Writer, ABC):
                     # get label hierarchy
                     # multiple labels:
                     if not force:
-                        all_labels = self.translator.ontology.get_ancestors(
-                            label
-                        )
+                        all_labels = self.translator.ontology.get_ancestors(label)
                     else:
                         all_labels = None
 
                     if all_labels:
                         # convert to pascal case
-                        all_labels = [
-                            self.translator.name_sentence_to_pascal(label)
-                            for label in all_labels
-                        ]
+                        all_labels = [self.translator.name_sentence_to_pascal(label) for label in all_labels]
                         # remove duplicates
                         all_labels = list(OrderedDict.fromkeys(all_labels))
                         # order alphabetically
@@ -502,9 +477,7 @@ class _BatchWriter(_Writer, ABC):
                         # concatenate with array delimiter
                         all_labels = self._write_array_string(all_labels)
                     else:
-                        all_labels = self.translator.name_sentence_to_pascal(
-                            label
-                        )
+                        all_labels = self.translator.name_sentence_to_pascal(label)
 
                     labels[label] = all_labels
 
@@ -549,16 +522,15 @@ class _BatchWriter(_Writer, ABC):
                 self.node_property_dict[label] = reference_props[label]
 
             return True
+        elif not isinstance(nodes, list):
+            logger.error("Nodes must be passed as list or generator.")
+            return False
         else:
-            if type(nodes) is not list:
-                logger.error("Nodes must be passed as list or generator.")
-                return False
-            else:
 
-                def gen(nodes):
-                    yield from nodes
+            def gen(nodes):
+                yield from nodes
 
-                return self._write_node_data(gen(nodes), batch_size=batch_size)
+            return self._write_node_data(gen(nodes), batch_size=batch_size)
 
     def _write_single_node_list_to_file(
         self,
@@ -567,11 +539,11 @@ class _BatchWriter(_Writer, ABC):
         prop_dict: dict,
         labels: str,
     ):
-        """
-        This function takes one list of biocypher nodes and writes them
+        """This function takes one list of biocypher nodes and writes them
         to a Neo4j admin import compatible CSV file.
 
         Args:
+        ----
             node_list (list): list of BioCypherNodes to be written
             label (str): the primary label of the node
             prop_dict (dict): properties of node class passed from parsing
@@ -580,7 +552,9 @@ class _BatchWriter(_Writer, ABC):
                 for the node class
 
         Returns:
+        -------
             bool: The return value. True for success, False otherwise.
+
         """
         if not all(isinstance(n, BioCypherNode) for n in node_list):
             logger.error("Nodes must be passed as type BioCypherNode.")
@@ -598,7 +572,7 @@ class _BatchWriter(_Writer, ABC):
             ref_props = list(prop_dict.keys())
 
             # compare lists order invariant
-            if not set(ref_props) == set(n_keys):
+            if set(ref_props) != set(n_keys):
                 onode = n.get_id()
                 oprop1 = set(ref_props).difference(n_keys)
                 oprop2 = set(n_keys).difference(ref_props)
@@ -632,11 +606,10 @@ class _BatchWriter(_Writer, ABC):
                         "boolean",
                     ]:
                         plist.append(str(p))
+                    elif isinstance(p, list):
+                        plist.append(self._write_array_string(p))
                     else:
-                        if isinstance(p, list):
-                            plist.append(self._write_array_string(p))
-                        else:
-                            plist.append(self._quote_string(str(p)))
+                        plist.append(f"{self.quote}{p!s}{self.quote}")
 
                 line.append(self.delim.join(plist))
             line.append(labels)
@@ -650,8 +623,7 @@ class _BatchWriter(_Writer, ABC):
         return True
 
     def _write_edge_data(self, edges, batch_size):
-        """
-        Writes biocypher edges to CSV conforming to the headers created
+        """Writes biocypher edges to CSV conforming to the headers created
         with `_write_edge_headers()`, and is actually required to be run
         before calling `_write_node_headers()` to set the
         :py:attr:`self.edge_property_dict` for passing the edge
@@ -659,17 +631,20 @@ class _BatchWriter(_Writer, ABC):
         from the :py:class:`BioCypherEdge` class.
 
         Args:
+        ----
             edges (BioCypherEdge): a list or generator of edges in
                 :py:class:`BioCypherEdge` format
 
         Returns:
+        -------
             bool: The return value. True for success, False otherwise.
 
         Todo:
+        ----
             - currently works for mixed edges but in practice often is
               called on one iterable containing one type of edge only
-        """
 
+        """
         if isinstance(edges, GeneratorType):
             logger.debug("Writing edge CSV from generator.")
 
@@ -685,14 +660,13 @@ class _BatchWriter(_Writer, ABC):
             for edge in edges:
                 if not (edge.get_source_id() and edge.get_target_id()):
                     logger.error(
-                        "Edge must have source and target node. "
-                        f"Caused by: {edge}",
+                        f"Edge must have source and target node. Caused by: {edge}",
                     )
                     continue
 
                 label = edge.get_label()
 
-                if not label in bins.keys():
+                if label not in bins.keys():
                     # start new list
                     bins[label].append(edge)
                     bin_l[label] = 1
@@ -703,13 +677,8 @@ class _BatchWriter(_Writer, ABC):
                     # (may not be if it is an edge that carries the
                     # "label_as_edge" property)
                     cprops = None
-                    if (
-                        label
-                        in self.translator.ontology.mapping.extended_schema
-                    ):
-                        cprops = self.translator.ontology.mapping.extended_schema.get(
-                            label
-                        ).get(
+                    if label in self.translator.ontology.mapping.extended_schema:
+                        cprops = self.translator.ontology.mapping.extended_schema.get(label).get(
                             "properties",
                         )
                     else:
@@ -717,9 +686,7 @@ class _BatchWriter(_Writer, ABC):
                         for (
                             k,
                             v,
-                        ) in (
-                            self.translator.ontology.mapping.extended_schema.items()
-                        ):
+                        ) in self.translator.ontology.mapping.extended_schema.items():
                             if isinstance(v, dict):
                                 if v.get("label_as_edge") == label:
                                     cprops = v.get("properties")
@@ -789,16 +756,15 @@ class _BatchWriter(_Writer, ABC):
                 self.edge_property_dict[label] = reference_props[label]
 
             return True
+        elif not isinstance(edges, list):
+            logger.error("Edges must be passed as list or generator.")
+            return False
         else:
-            if type(edges) is not list:
-                logger.error("Edges must be passed as list or generator.")
-                return False
-            else:
 
-                def gen(edges):
-                    yield from edges
+            def gen(edges):
+                yield from edges
 
-                return self._write_edge_data(gen(edges), batch_size=batch_size)
+            return self._write_edge_data(gen(edges), batch_size=batch_size)
 
     def _write_single_edge_list_to_file(
         self,
@@ -806,11 +772,11 @@ class _BatchWriter(_Writer, ABC):
         label: str,
         prop_dict: dict,
     ):
-        """
-        This function takes one list of biocypher edges and writes them
+        """This function takes one list of biocypher edges and writes them
         to a Neo4j admin import compatible CSV file.
 
         Args:
+        ----
             edge_list (list): list of BioCypherEdges to be written
 
             label (str): the label (type) of the edge
@@ -819,9 +785,10 @@ class _BatchWriter(_Writer, ABC):
                 function and their types
 
         Returns:
+        -------
             bool: The return value. True for success, False otherwise.
-        """
 
+        """
         if not all(isinstance(n, BioCypherEdge) for n in edge_list):
             logger.error("Edges must be passed as type BioCypherEdge.")
             return False
@@ -836,7 +803,7 @@ class _BatchWriter(_Writer, ABC):
             ref_props = list(prop_dict.keys())
 
             # compare list order invariant
-            if not set(ref_props) == set(e_keys):
+            if set(ref_props) != set(e_keys):
                 oedge = f"{e.get_source_id()}-{e.get_target_id()}"
                 oprop1 = set(ref_props).difference(e_keys)
                 oprop2 = set(e_keys).difference(ref_props)
@@ -867,11 +834,10 @@ class _BatchWriter(_Writer, ABC):
                     "boolean",
                 ]:
                     plist.append(str(p))
+                elif isinstance(p, list):
+                    plist.append(self._write_array_string(p))
                 else:
-                    if isinstance(p, list):
-                        plist.append(self._write_array_string(p))
-                    else:
-                        plist.append(self._quote_string(str(p)))
+                    plist.append(self.quote + str(p) + self.quote)
 
             entries = [e.get_source_id()]
 
@@ -880,9 +846,7 @@ class _BatchWriter(_Writer, ABC):
 
             if label in ["IS_SOURCE_OF", "IS_TARGET_OF", "IS_PART_OF"]:
                 skip_id = True
-            elif not self.translator.ontology.mapping.extended_schema.get(
-                label
-            ):
+            elif not self.translator.ontology.mapping.extended_schema.get(label):
                 # find label in schema by label_as_edge
                 for (
                     k,
@@ -897,9 +861,9 @@ class _BatchWriter(_Writer, ABC):
             if schema_label:
                 if (
                     self.translator.ontology.mapping.extended_schema.get(
-                        schema_label
+                        schema_label,
                     ).get("use_id")
-                    == False
+                    == False  # noqa: E712 (seems to not work with 'not')
                 ):
                     skip_id = True
 
@@ -913,7 +877,7 @@ class _BatchWriter(_Writer, ABC):
             entries.append(
                 self.translator.name_sentence_to_pascal(
                     e.get_label(),
-                )
+                ),
             )
 
             lines.append(
@@ -927,10 +891,10 @@ class _BatchWriter(_Writer, ABC):
         return True
 
     def _write_next_part(self, label: str, lines: list):
-        """
-        This function writes a list of strings to a new part file.
+        """This function writes a list of strings to a new part file.
 
         Args:
+        ----
             label (str): the label (type) of the edge; internal
             representation sentence case -> needs to become PascalCase
             for disk representation
@@ -938,17 +902,15 @@ class _BatchWriter(_Writer, ABC):
             lines (list): list of strings to be written
 
         Returns:
+        -------
             bool: The return value. True for success, False otherwise.
+
         """
         # translate label to PascalCase
-        label_pascal = self.translator.name_sentence_to_pascal(
-            parse_label(label)
-        )
+        label_pascal = self.translator.name_sentence_to_pascal(parse_label(label))
 
         # list files in self.outdir
-        files = glob.glob(
-            os.path.join(self.outdir, f"{label_pascal}-part*.csv")
-        )
+        files = glob.glob(os.path.join(self.outdir, f"{label_pascal}-part*.csv"))
         # find file with highest part number
         if not files:
             next_part = 0
@@ -956,10 +918,7 @@ class _BatchWriter(_Writer, ABC):
         else:
             next_part = (
                 max(
-                    [
-                        int(f.split(".")[-2].split("-")[-1].replace("part", ""))
-                        for f in files
-                    ],
+                    [int(f.split(".")[-2].split("-")[-1].replace("part", "")) for f in files],
                 )
                 + 1
             )
@@ -984,31 +943,29 @@ class _BatchWriter(_Writer, ABC):
             self.parts[label].append(part)
 
     def get_import_call(self) -> str:
-        """
-        Function to return the import call detailing folder and
+        """Function to return the import call detailing folder and
         individual node and edge headers and data files, as well as
         delimiters and database name.
 
-        Returns:
+        Returns
+        -------
             str: a bash command for the database import
-        """
 
+        """
         return self._construct_import_call()
 
     def write_import_call(self) -> str:
-        """
-        Function to write the import call detailing folder and
+        """Function to 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:
+        Returns
+        -------
             str: The path of the file holding the import call.
-        """
 
+        """
         file_path = os.path.join(self.outdir, self._get_import_script_name())
-        logger.info(
-            f"Writing {self.db_name + ' ' if self.db_name else ''}import call to `{file_path}`."
-        )
+        logger.info(f"Writing {self.db_name + ' ' if self.db_name else ''}import call to `{file_path}`.")
 
         with open(file_path, "w", encoding="utf-8") as f:
             f.write(self._construct_import_call())
@@ -1017,16 +974,16 @@ class _BatchWriter(_Writer, ABC):
 
 
 def parse_label(label: str) -> str:
-    """
-
-    Check if the label is compliant with Neo4j naming conventions,
+    """Check if the label is compliant with Neo4j naming conventions,
     https://neo4j.com/docs/cypher-manual/current/syntax/naming/, and if not,
     remove non-compliant characters.
 
     Args:
+    ----
         label (str): The label to check
     Returns:
         str: The compliant label
+
     """
     # Check if the name contains only alphanumeric characters, underscore, or dollar sign
     # and dot (for class hierarchy of BioCypher)
@@ -1036,7 +993,7 @@ def parse_label(label: str) -> str:
     if non_matches:
         non_matches = list(set(non_matches))
         logger.warning(
-            f"Label is not compliant with Neo4j naming rules. Removed non compliant characters: {non_matches}"
+            f"Label is not compliant with Neo4j naming rules. Removed non compliant characters: {non_matches}",
         )
 
     def first_character_compliant(character: str) -> bool:
@@ -1047,7 +1004,5 @@ def parse_label(label: str) -> str:
             if first_character_compliant(c):
                 matches = matches[matches.index(c) :]
                 break
-        logger.warning(
-            "Label does not start with an alphabetic character or with $. Removed non compliant characters."
-        )
+        logger.warning("Label does not start with an alphabetic character or with $. Removed non compliant characters.")
     return "".join(matches).strip()