Class HeteroEdgeDataView

python/dgl/view.py:182–267

The data view class when G.edata[etype] is called.

class HeteroEdgeDataView(MutableMapping):
    """The data view class when G.edata[etype] is called."""

    __slots__ = ["_graph", "_etype", "_etid", "_edges"]

    def __init__(self, graph, etype, edges):
        self._graph = graph
        self._etype = etype
        self._etid = (
            [self._graph.get_etype_id(t) for t in etype]
            if isinstance(etype, list)
            else self._graph.get_etype_id(etype)
        )
        self._edges = edges

    def __getitem__(self, key):
        if isinstance(self._etype, list):
            ret = {}
            for (i, etype) in enumerate(self._etype):
                value = self._graph._get_e_repr(self._etid[i], self._edges).get(
                    key, None
                )
                if value is not None:
                    ret[etype] = value
            return ret
        else:
            return self._graph._get_e_repr(self._etid, self._edges)[key]

    def __setitem__(self, key, val):
        if isinstance(val, LazyFeature):
            self._graph._edge_frames[self._etid][key] = val
        elif isinstance(self._etype, list):
            assert isinstance(val, dict), (
                "Current HeteroEdgeDataView has multiple edge types, "
                "please pass the edge type and the corresponding data through a dict."
            )

            for (etype, data) in val.items():
                etid = self._graph.get_etype_id(etype)
                self._graph._set_e_repr(etid, self._edges, {key: data})
        else:
            assert isinstance(val, dict) is False, (
                "The HeteroEdgeDataView has only one edge type. "
                "please pass a tensor directly"
            )
            self._graph._set_e_repr(self._etid, self._edges, {key: val})

    def __delitem__(self, key):
        if isinstance(self._etype, list):
            for etid in self._etid:
                if self._graph._get_e_repr(etid, ALL).get(key, None) is None:
                    continue
                self._graph._pop_e_repr(etid, key)
        else:
            self._graph._pop_e_repr(self._etid, key)

    def _transpose(self, as_dict=False):
        if isinstance(self._etype, list):