overlay.py

class NodeDictOverlay(MutableMapping):
    def __init__(self, frame):
        self._frame = frame

    @property
    def num_nodes(self):
        return self._frame.num_rows()

    def add_nodes(self, nodes, attrs):
        # NOTE: currently `nodes` are not used. Users need to make sure
        # the node ids are continuous ids from 0.
        # NOTE: this is a good place to hook any graph mutation logic.
        self._frame.append(attrs)

    def delete_nodes(self, nodes):
        # NOTE: this is a good place to hook any graph mutation logic.
        raise NotImplementedError('Delete nodes in the graph is currently not supported.')

    def get_node_attrs(self, nodes, key):
        if nodes == ALL:
            # get the whole column
            return self._frame[key]
        else:
            # TODO(minjie): should not rely on tensor's __getitem__ syntax.
            return utils.id_type_dispatch(
                    nodes,
                    lambda nid : self._frame[key][nid],
                    lambda id_array : self._frame[key][id_array])

    def set_node_attrs(self, nodes, key, val):
        if nodes == ALL:
            # replace the whole column
            self._frame[key] = val
        else:
            # TODO(minjie): should not rely on tensor's __setitem__ syntax.
            utils.id_type_dispatch(
                    nodes,
                    lambda nid : self._frame[key][nid] = val,
                    lambda id_array : self._frame[key][id_array] = val)

    def __getitem__(self, nodes):
        def _check_one(nid):
            if nid >= self.num_nodes:
                raise KeyError
        def _check_many(id_array):
            if F.max(id_array) >= self.num_nodes:
                raise KeyError
        utils.id_type_dispatch(nodes, _check_one, _check_many)
        return utils.MutableLazyDict(
                lambda key: self.get_node_attrs(nodes, key),
                lambda key, val: self.set_node_attrs(nodes, key, val)
                self._frame.schemes)

    def __setitem__(self, nodes, attrs):
        # Happens when adding new nodes in the graph.
        self.add_nodes(nodes, attrs)

    def __delitem__(self, nodes):
        # Happens when deleting nodes in the graph.
        self.delete_nodes(nodes)

    def __len__(self):
        return self.num_nodes

    def __iter__(self):
        raise NotImplementedError()

class AdjOuterOverlay(MutableMapping):
    """
    TODO: Replace this with a more efficient dict structure.
    TODO: Batch graph mutation is not supported.
    """
    def __init__(self):
        self._adj = {}

    def __setitem__(self, u, inner_dict):
        self._adj[u] = inner_dict

    def __getitem__(self, u):
        def _check_one(nid):
            if nid not in self._adj:
                raise KeyError
        def _check_many(id_array):
            pass
        utils.id_type_dispatch(u, _check_one, _check_many)
        return utils.id_type_dispatch(u)

    def __delitem__(self, u):
        # The delitem is ignored.
        raise NotImplementedError('Delete edges in the graph is currently not supported.')

class AdjInnerOverlay(dict):
    """TODO: replace this with a more efficient dict structure."""
    def __setitem__(self, v, attrs):
        pass