<commit_before>import networkx as nx
import pandas as pd


class PandasGraph(object):
    def __init__(self, data):
        self.__data = data.sortlevel(level=0)
        self.__dg = nx.DiGraph()

        for ids in data.index:
            self.__dg.add_node(ids[0])
            self.__dg.add_node(ids[1])
            self.__dg.add_edge(ids[0], ids[1])

    @property
    def data(self):
        return self.__data

    def __edge(self, owner, owns):
        return self.__data.xs([owner, owns], level=["Owner ID", "Owned ID"], drop_level=False)

    def __edges(self, id):
        """ List of all edges reachable from node id """
        return pd.concat([self.__edge(owner=e[0], owns=e[1]) for e in nx.dfs_edges(self.__dg, id)], axis=0)

    def left(self, ids):
        assert isinstance(ids, list)
        # all edges reachable from a list of nodes
        return pd.concat([self.__edges(id) for id in ids], axis=0)<commit_msg>make pandasgraph robuster<commit_after>import networkx as nx
import pandas as pd


class PandasGraph(object):
    """
    Given a Pandas Dataframe, we can take two index-columns called left and right to create a graph G with edges from
    the 'left' to the 'right'.
    """
    def __init__(self, data, left="Owner ID", right="Owned ID"):
        """
        We create a graph based on two columns in the MultiIndex of a Pandas DataFrame.

        :param data: A pandas DataFrame with a MultiIndex.
        :param left: The name of the level for the starting nodes (e.g. we start all edges in them)
        :param right: The name of the level for the ending nodes (e.g. we end all edges in them)
        """
        self.__data = data.sortlevel(level=0)
        self.__dg = nx.DiGraph()
        self.__left = left
        self.__right = right

        graph = zip(data.index.get_level_values(level=left), data.index.get_level_values(level=right))

        for left_n, right_n in graph:
            self.__dg.add_node(left_n)
            self.__dg.add_node(right_n)
            self.__dg.add_edge(left_n, right_n)

    @property
    def data(self):
        return self.__data

    def __edge(self, left, right):
        return self.__data.xs([left, right], level=[self.__left, self.__right], drop_level=False)

    def edges(self, ids):
        """
        Return all edges reachable from a node described in the list ids.
        If A owns B and B owns C, we get[[A,B],[B,C]]
        """
        assert isinstance(ids, list)
        # all edges reachable from a list of nodes
        return pd.concat([pd.concat([self.__edge(e[0], e[1]) for e in nx.dfs_edges(self.__dg, id)]) for id in ids])
{'fraction_non_alphanumeric': 0.11174675819984745, 'fraction_numerical': 0.004576659038901602, 'mean_word_length': 2.563858695652174, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6515312', 'n_tokens_mistral': 834, 'n_tokens_neox': 804, 'n_words': 269}