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

biocypher/output/write/graph/_networkx.py

@@ -1,31 +1,34 @@
 import pickle
 
-import networkx as nx
-
 from biocypher._logger import logger
+from biocypher.output.in_memory._networkx import NetworkxKG
 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.
+    Class for writing the in-memory networkx DiGraph to file.
+
+    Call `_construct_import_call` to write the networkx DiGraph to a pickle
+    file and return the Python call to load it.
+
+    TODO: this is a non-intuitive name, should be adjusted.
     """
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.csv_writer = _PandasCSVWriter(*args, write_to_file=False, **kwargs)
-        self.G = nx.DiGraph()
+        self.in_memory_networkx_kg = NetworkxKG(
+            deduplicator=self.deduplicator,
+        )
 
     def _construct_import_call(self) -> str:
-        """Function to construct the Python code to load all node and edge csv files again into Pandas dfs.
+        """Dump networkx graph to a pickle file and return Python call.
 
         Returns:
-            str: Python code to load the csv files into Pandas dfs.
+            str: Python code to load the networkx graph from a pickle file.
         """
-        logger.info(
-            f"Writing networkx {self.G} to pickle file networkx_graph.pkl."
-        )
+        self.G = self.in_memory_networkx_kg._create_networkx_kg()
+        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)
 
@@ -38,39 +41,29 @@ class _NetworkXWriter(_Writer):
         return "import_networkx.py"
 
     def _write_node_data(self, nodes) -> bool:
-        passed = self.csv_writer._write_entities_to_file(nodes)
-        self.add_to_networkx()
+        """Add nodes to the networkx graph.
+
+        TODO: this is not strictly writing, should be refactored.
+
+        Args:
+            nodes (list): List of nodes to add to the networkx graph.
+
+        Returns:
+            bool: True if the nodes were added successfully, False otherwise.
+        """
+        passed = self.in_memory_networkx_kg.add_nodes(nodes)
         return passed
 
     def _write_edge_data(self, edges) -> bool:
-        passed = self.csv_writer._write_entities_to_file(edges)
-        self.add_to_networkx()
-        return passed
+        """Add edges to the networkx graph.
+
+        TODO: this is not strictly writing, should be refactored.
 
-    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
+        Args:
+            edges (list): List of edges to add to the networkx graph.
+
+        Returns:
+            bool: True if the edges were added successfully, False otherwise.
+        """
+        passed = self.in_memory_networkx_kg.add_edges(edges)
+        return passed

biocypher/output/write/graph/_rdf.py

@@ -1,22 +1,12 @@
-#!/usr/bin/env python
-
-#
-# Copyright 2021, Heidelberg University Clinic
-#
-# File author(s):  Loes van den Biggelaar
-#                  Sebastian Lobentanzer
-#
-# Distributed under MIT licence, see the file `LICENSE`.
-#
-"""
-BioCypher 'offline' module. Handles the writing of node and edge representations
+"""BioCypher 'offline' module. Handles the writing of node and edge representations
 suitable for import into a DBMS.
 """
-from types import GeneratorType
-from typing import Union
+
 import os
 
-from rdflib import DC, RDF, RDFS, SKOS, DCTERMS, Graph, Literal, Namespace
+from types import GeneratorType
+
+from rdflib import DC, DCTERMS, RDF, RDFS, SKOS, Graph, Literal, Namespace
 from rdflib.namespace import (
     _NAMESPACE_PREFIXES_CORE,
     _NAMESPACE_PREFIXES_RDFLIB,
@@ -28,8 +18,7 @@ from biocypher.output.write._batch_writer import _BatchWriter
 
 
 class _RDFWriter(_BatchWriter):
-    """
-    Class to write BioCypher's property graph into an RDF format using
+    """Class to write BioCypher's property graph into an RDF format using
     rdflib and all the extensions it supports (RDF/XML, N3, NTriples,
     N-Quads, Turtle, TriX, Trig and JSON-LD). By default the conversion
     is done keeping only the minimum information about node and edges,
@@ -37,33 +26,37 @@ class _RDFWriter(_BatchWriter):
     """
 
     def _get_import_script_name(self) -> str:
-        """
-        Returns the name of the RDF admin import script.
+        """Returns the name of the RDF admin import script.
         This function applicable for RDF export.
 
-        Returns:
+        Returns
+        -------
             str: The name of the import script (ending in .sh)
+
         """
         return "rdf-import-call.sh"
 
     def _get_default_import_call_bin_prefix(self):
-        """
-        Method to provide the default string for the import call bin prefix.
+        """Method to provide the default string for the import call bin prefix.
 
-        Returns:
+        Returns
+        -------
             str: The default location for the RDF admin import location
+
         """
         return "bin/"
 
     def _is_rdf_format_supported(self, rdf_format: str) -> bool:
-        """
-        Function to check if the specified RDF format is supported.
+        """Function to check if the specified RDF format is supported.
 
         Args:
+        ----
             rdf_format (str): The RDF format to check.
 
         Returns:
+        -------
             bool: Returns True if rdf format supported, False otherwise.
+
         """
         supported_formats = [
             "xml",
@@ -83,7 +76,8 @@ class _RDFWriter(_BatchWriter):
             )
             return False
         else:
-            # RDF graph does not support 'ttl' format, only 'turtle' format. however, the preferred file extension is always '.ttl'
+            # RDF graph does not support 'ttl' format, only 'turtle' format.
+            # however, the preferred file extension is always '.ttl'
             if self.rdf_format == "turtle":
                 self.extension = "ttl"
             elif self.rdf_format == "ttl":
@@ -99,11 +93,11 @@ class _RDFWriter(_BatchWriter):
         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 an RDF file with the given format.
 
         Args:
+        ----
             edge_list (list): list of BioCypherEdges to be written
 
             label (str): the label (type) of the edge
@@ -112,9 +106,10 @@ class _RDFWriter(_BatchWriter):
                 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
@@ -123,9 +118,7 @@ class _RDFWriter(_BatchWriter):
         label_pascal = self.translator.name_sentence_to_pascal(label)
 
         # create file name
-        file_name = os.path.join(
-            self.outdir, f"{label_pascal}.{self.extension}"
-        )
+        file_name = os.path.join(self.outdir, f"{label_pascal}.{self.extension}")
 
         # write data in graph
         graph = Graph()
@@ -136,12 +129,10 @@ class _RDFWriter(_BatchWriter):
             rdf_object = edge.get_target_id()
             rdf_predicate = edge.get_id()
             rdf_properties = edge.get_properties()
-            if rdf_predicate == None:
+            if rdf_predicate is None:
                 rdf_predicate = rdf_subject + rdf_object
 
-            edge_label = self.translator.name_sentence_to_pascal(
-                edge.get_label()
-            )
+            edge_label = self.translator.name_sentence_to_pascal(edge.get_label())
             edge_uri = self.rdf_namespaces["biocypher"][edge_label]
             graph.add((edge_uri, RDF.type, RDFS.Class))
             graph.add(
@@ -149,21 +140,21 @@ class _RDFWriter(_BatchWriter):
                     self.rdf_namespaces["biocypher"][rdf_predicate],
                     RDF.type,
                     edge_uri,
-                )
+                ),
             )
             graph.add(
                 (
                     self.rdf_namespaces["biocypher"][rdf_predicate],
                     self.rdf_namespaces["biocypher"]["subject"],
                     self.subject_to_uri(rdf_subject),
-                )
+                ),
             )
             graph.add(
                 (
                     self.rdf_namespaces["biocypher"][rdf_predicate],
                     self.rdf_namespaces["biocypher"]["object"],
                     self.subject_to_uri(rdf_object),
-                )
+                ),
             )
 
             # add properties to the transformed edge --> node
@@ -187,13 +178,17 @@ class _RDFWriter(_BatchWriter):
         rdf_object: str,
         rdf_predicate: str,
     ):
-        """
-        Function to add the properties to an RDF node. It takes the graph, the subject, object, and predicate of the RDF triple.
-        It checks if the property is a list and adds it to the graph accordingly. otherwise it checks if the string represents a list.
-        If it does, it transforms it to a list and adds it to the graph. if not, it adds the property to the graph as a literal.
-        If the property is neither a list or string, it will also be added as a literal.
+        """Add the properties to an RDF node.
+
+        It takes the graph, the subject, object, and predicate of the RDF
+        triple. It checks if the property is a list and adds it to the graph
+        accordingly. Otherwise it checks if the string represents a list. If it
+        does, it transforms it to a list and adds it to the graph. If not, it
+        adds the property to the graph as a literal. If the property is neither
+        a list or string, it will also be added as a literal.
 
         Args:
+        ----
             graph (RDFLib.Graph): The RDF graph to add the nodes to.
 
             rdf_subject (str): The subject of the RDF triple.
@@ -203,7 +198,9 @@ class _RDFWriter(_BatchWriter):
             rdf_predicate (str): The predicate of the RDF triple.
 
         Returns:
+        -------
             None
+
         """
         if isinstance(rdf_object, list):
             for obj in rdf_object:
@@ -212,7 +209,7 @@ class _RDFWriter(_BatchWriter):
                         self.subject_to_uri(rdf_subject),
                         self.property_to_uri(rdf_predicate),
                         Literal(obj),
-                    )
+                    ),
                 )
         elif isinstance(rdf_object, str):
             if rdf_object.startswith("[") and rdf_object.endswith("]"):
@@ -228,7 +225,7 @@ class _RDFWriter(_BatchWriter):
                         self.subject_to_uri(rdf_subject),
                         self.property_to_uri(rdf_predicate),
                         Literal(rdf_object),
-                    )
+                    ),
                 )
         else:
             graph.add(
@@ -236,25 +233,22 @@ class _RDFWriter(_BatchWriter):
                     self.subject_to_uri(rdf_subject),
                     self.property_to_uri(rdf_predicate),
                     Literal(rdf_object),
-                )
+                ),
             )
 
     def transform_string_to_list(self, string_list: str) -> list:
-        """
-        Function to transform a string representation of a list into a list.
+        """Function to transform a string representation of a list into a list.
 
         Args:
+        ----
             string_list (str): The string representation of the list.
 
         Returns:
+        -------
             list: The list representation of the input string.
+
         """
-        return (
-            string_list.replace("[", "")
-            .replace("]", "")
-            .replace("'", "")
-            .split(", ")
-        )
+        return string_list.replace("[", "").replace("]", "").replace("'", "").split(", ")
 
     def _write_single_node_list_to_file(
         self,
@@ -263,11 +257,11 @@ class _RDFWriter(_BatchWriter):
         prop_dict: dict,
         labels: str,
     ):
-        """
-        This function takes a list of BioCypherNodes and writes them
+        """This function takes a list of BioCypherNodes and writes them
         to an RDF file in the specified format.
 
         Args:
+        ----
             node_list (list): A list of BioCypherNodes to be written.
 
             label (str): The label (type) of the nodes.
@@ -275,7 +269,9 @@ class _RDFWriter(_BatchWriter):
             prop_dict (dict): A dictionary of properties and their types for the node class.
 
         Returns:
+        -------
             bool: True if the writing is successful, False otherwise.
+
         """
         if not all(isinstance(n, BioCypherNode) for n in node_list):
             logger.error("Nodes must be passed as type BioCypherNode.")
@@ -285,9 +281,7 @@ class _RDFWriter(_BatchWriter):
         label_pascal = self.translator.name_sentence_to_pascal(label)
 
         # create file name
-        file_name = os.path.join(
-            self.outdir, f"{label_pascal}.{self.extension}"
-        )
+        file_name = os.path.join(self.outdir, f"{label_pascal}.{self.extension}")
 
         # write data in graph
         graph = Graph()
@@ -303,14 +297,14 @@ class _RDFWriter(_BatchWriter):
                     self.rdf_namespaces["biocypher"][class_name],
                     RDF.type,
                     RDFS.Class,
-                )
+                ),
             )
             graph.add(
                 (
                     self.subject_to_uri(rdf_subject),
                     RDF.type,
                     self.rdf_namespaces["biocypher"][class_name],
-                )
+                ),
             )
             for key, value in properties.items():
                 # only write value if it exists.
@@ -325,19 +319,19 @@ class _RDFWriter(_BatchWriter):
 
         return True
 
-    def write_nodes(
-        self, nodes, batch_size: int = int(1e6), force: bool = False
-    ) -> bool:
-        """
-        Wrapper for writing nodes in RDF format. It calls the _write_node_data() function, specifying the node data.
+    def write_nodes(self, nodes, batch_size: int = int(1e6), force: bool = False) -> bool:
+        """Wrapper for writing nodes in RDF format. It calls the _write_node_data() function, specifying the node data.
 
         Args:
+        ----
             nodes (list or generator): A list or generator of nodes in BioCypherNode format.
             batch_size (int): The number of nodes to write in each batch.
             force (bool): Flag to force the writing even if the output file already exists.
 
         Returns:
+        -------
             bool: True if the writing is successful, False otherwise.
+
         """
         # check if specified output format is correct
         passed = self._is_rdf_format_supported(self.rdf_format)
@@ -353,20 +347,22 @@ class _RDFWriter(_BatchWriter):
 
     def write_edges(
         self,
-        edges: Union[list, GeneratorType],
+        edges: list | GeneratorType,
         batch_size: int = int(1e6),
     ) -> bool:
-        """
-        Wrapper for writing edges in RDF format. It calls _write_edge_data()
+        """Wrapper for writing edges in RDF format. It calls _write_edge_data()
         functions specifying it's edge data.
 
         Args:
+        ----
             edges (BioCypherEdge): a list or generator of edges in
                 :py:class:`BioCypherEdge` format
             batch_size (int): The number of edges to write in each batch.
 
         Returns:
+        -------
             bool: The return value. True for success, False otherwise.
+
         """
         # check if specified output format is correct
         passed = self._is_rdf_format_supported(self.rdf_format)
@@ -382,12 +378,13 @@ class _RDFWriter(_BatchWriter):
         return True
 
     def _construct_import_call(self) -> bool:
-        """
-        Function to write the import call.
+        """Function to write the import call.
         This function is not applicable for RDF.
 
-        Returns:
+        Returns
+        -------
             bool: The return value. True for success, False otherwise.
+
         """
         return ""
 
@@ -399,53 +396,58 @@ class _RDFWriter(_BatchWriter):
         return f"{self.quote}{value}{self.quote}"
 
     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
         as required by the RDF admin-import.
         This function is not applicable for RDF.
 
         Args:
+        ----
             string_list (list): list of ontology strings
 
         Returns:
+        -------
             str: The string representation of an array for the neo4j admin import
-        """
 
+        """
         return True
 
     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`
         This function is not applicable for RDF.
 
-        Returns:
+        Returns
+        -------
             bool: The return value. True for success, False otherwise.
+
         """
         return True
 
     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.
         This function is not applicable for RDF.
 
-        Returns:
+        Returns
+        -------
             bool: The return value. True for success, False otherwise.
+
         """
         return True
 
     def subject_to_uri(self, subject: str) -> str:
-        """
-        Converts the subject to a proper URI using the available namespaces.
+        """Converts the subject to a proper URI using the available namespaces.
         If the conversion fails, it defaults to the biocypher prefix.
 
         Args:
+        ----
             subject (str): The subject to be converted to a URI.
 
         Returns:
+        -------
             str: The corresponding URI for the subject.
+
         """
         try:
             _pref, _id = subject.split(":")
@@ -458,56 +460,66 @@ class _RDFWriter(_BatchWriter):
             return self.rdf_namespaces["biocypher"][subject]
 
     def property_to_uri(self, property_name: str) -> dict[str, str]:
-        """
-        Converts a property name to its corresponding URI.
+        """Converts a property name to its corresponding URI.
 
         This function takes a property name and searches for its corresponding URI in various namespaces.
         It first checks the core namespaces for rdflib, including owl, rdf, rdfs, xsd, and xml.
 
         Args:
+        ----
             property_name (str): The property name to be converted to a URI.
 
         Returns:
+        -------
             str: The corresponding URI for the input property name.
+
         """
         # These namespaces are core for rdflib; owl, rdf, rdfs, xsd and xml
         for namespace in _NAMESPACE_PREFIXES_CORE.values():
             if property_name in namespace:
                 return namespace[property_name]
 
-        # If the property name is not found in the core namespaces, search in the SKOS, DC, and DCTERMS namespaces
+        # If the property name is not found in the core namespaces, search in
+        # the SKOS, DC, and DCTERMS namespaces
         for namespace in [SKOS, DC, DCTERMS]:
             if property_name in namespace:
                 return namespace[property_name]
 
-        # If the property name is still not found, try other namespaces from rdflib.
+        # If the property name is still not found, try other namespaces from
+        # rdflib.
         for namespace in _NAMESPACE_PREFIXES_RDFLIB.values():
             if property_name in namespace:
                 return namespace[property_name]
 
-        # If the property name is "licence", it recursively calls the function with "license" as the input.
+        # If the property name is "licence", it recursively calls the function
+        # with "license" as the input.
         if property_name == "licence":
             return self.property_to_uri("license")
 
         # TODO: add an option to search trough manually implemented namespaces
 
-        # If the input is not found in any of the namespaces, it returns the corresponding URI from the biocypher namespace.
+        # If the input is not found in any of the namespaces, it returns
+        # the corresponding URI from the biocypher namespace.
         # TODO: give a warning and try to prevent this option altogether
         return self.rdf_namespaces["biocypher"][property_name]
 
     def _init_namespaces(self, graph: Graph):
-        """
-        Initializes the namespaces for the RDF graph. These namespaces are used to convert nodes to URIs.
+        """Initialise the namespaces for the RDF graph.
 
-        This function adds the biocypher standard namespace to the `rdf_namespaces` attribute of the class.
-        If `rdf_namespaces` is empty, it sets it to the biocypher standard namespace. Otherwise, it merges
-        the biocypher standard namespace with the namespaces defined in the biocypher_config.yaml.
+        These namespaces are used to convert nodes to URIs. This function adds
+        the biocypher standard namespace to the `rdf_namespaces` attribute of
+        the class. If `rdf_namespaces` is empty, it sets it to the biocypher
+        standard namespace. Otherwise, it merges the biocypher standard
+        namespace with the namespaces defined in the biocypher_config.yaml.
 
         Args:
+        ----
             graph (RDFLib.Graph): The RDF graph to bind the namespaces to.
 
         Returns:
+        -------
             None
+
         """
         # add biocypher standard to self.rdf_namespaces
         biocypher_standard = {"biocypher": "https://biocypher.org/biocypher#"}

biocypher/output/write/relational/_csv.py

@@ -1,13 +1,13 @@
 from more_itertools import peekable
 
 from biocypher._logger import logger
+from biocypher.output.in_memory._pandas import PandasKG
 from biocypher.output.write._writer import _Writer
-from biocypher.output.in_memory._pandas import Pandas
 
 
 class _PandasCSVWriter(_Writer):
     """
-    Class for writing node and edge representations to a CSV file.
+    Class for writing node and edge representations to CSV files.
     """
 
     def __init__(self, *args, write_to_file: bool = True, **kwargs):
@@ -15,8 +15,7 @@ class _PandasCSVWriter(_Writer):
         super().__init__(*args, **kwargs)
         self.in_memory_dfs = {}
         self.stored_dfs = {}
-        self.pandas_in_memory = Pandas(
-            translator=self.translator,
+        self.pandas_in_memory = PandasKG(
             deduplicator=self.deduplicator,
         )
         self.delimiter = kwargs.get("delimiter")
@@ -48,7 +47,7 @@ class _PandasCSVWriter(_Writer):
         return passed
 
     def _write_entities_to_file(self, entities: iter) -> bool:
-        """Function to output.write the entities to a CSV file.
+        """Function to write the entities to a CSV file.
 
         Args:
             entities (iterable): An iterable of BioCypherNode / BioCypherEdge / BioCypherRelAsNode objects.
@@ -56,17 +55,13 @@ class _PandasCSVWriter(_Writer):
         entities = peekable(entities)
         entity_list = self.pandas_in_memory._separate_entity_types(entities)
         for entity_type, entities in entity_list.items():
-            self.in_memory_dfs[
-                entity_type
-            ] = self.pandas_in_memory._add_entity_df(entity_type, entities)
+            self.in_memory_dfs[entity_type] = self.pandas_in_memory._add_entity_df(entity_type, entities)
         for entity_type in self.in_memory_dfs.keys():
             entity_df = self.in_memory_dfs[entity_type]
             if " " in entity_type or "." in entity_type:
                 entity_type = entity_type.replace(" ", "_").replace(".", "_")
             if self.write_to_file:
-                logger.info(
-                    f"Writing {entity_df.shape[0]} entries to {entity_type}.csv."
-                )
+                logger.info(f"Writing {entity_df.shape[0]} entries to {entity_type}.csv.")
                 entity_df.to_csv(
                     f"{self.output_directory}/{entity_type}.csv",
                     sep=self.delimiter,