Source code for cfpq_data.graphs.readwrite.csv

"""Read (and write) a graph from (and to) CSV file."""
import logging
import pathlib
from typing import Union

import networkx as nx
import pandas as pd

__all__ = [
    "graph_from_csv",
    "graph_to_csv",
]




def graph_from_csv(path: Union[pathlib.Path, str]) -> nx.MultiDiGraph:
    """Loads a graph from CSV file.

    Parameters
    ----------
    path : Union[Path, str]
        The path to the CSV file with which
        the graph will be created.

    Examples
    --------
    >>> from cfpq_data import *
    >>> import pathlib, tempfile
    >>> p = pathlib.Path(tempfile.mkdtemp()) / "g.csv"
    >>> _ = p.write_text("0 1 a\\n1 2 b\\n")
    >>> g = graph_from_csv(p)
    >>> g.number_of_nodes()
    3
    >>> g.number_of_edges()
    2

    Returns
    -------
    g : MultiDiGraph
        Loaded graph.
    """
    data = pd.read_csv(
        filepath_or_buffer=path,
        sep=" ",
        header=None,
        names=["from", "to", "label"],
        engine="c",
    )

    graph = nx.from_pandas_edgelist(
        df=data,
        source="from",
        target="to",
        edge_attr="label",
        create_using=nx.MultiDiGraph,
    )

    logging.info(f"Load {graph=} from {path=}")

    return graph





def graph_to_csv(
    graph: nx.MultiDiGraph, path: Union[pathlib.Path, str]
) -> pathlib.Path:
    """Saves the ``graph`` to the CSV file by ``path``.

    Parameters
    ----------
    graph : MultiDiGraph
        Graph to save.

    path : Union[Path, str]
        The path to the CSV file where the graph will be saved.

    Examples
    --------
    >>> from cfpq_data import *
    >>> import pathlib, tempfile
    >>> d = pathlib.Path(tempfile.mkdtemp())
    >>> p = d / "g.csv"
    >>> _ = p.write_text("0 1 a\\n1 2 b\\n")
    >>> g = graph_from_csv(p)
    >>> path = graph_to_csv(g, d / "out.csv")
    >>> g2 = graph_from_csv(path)
    >>> sorted((u, v, e["label"]) for u, v, e in g2.edges(data=True)) == \\
    ...     sorted((u, v, e["label"]) for u, v, e in g.edges(data=True))
    True

    Returns
    -------
    path : Path
        Path to the CSV file where the graph will be saved.
    """
    with open(file=path, mode="w") as f:
        for u, v, edge_labels in graph.edges(data=True):
            for e in edge_labels.values():
                f.write(f"{u} {v} {e}\n")

    dest = pathlib.Path(path).resolve()

    logging.info(f"Save {graph=} to {dest=}")

    return dest