A NodeView class to act as G.nodes for a DGLGraph.
| 13 | |
| 14 | |
| 15 | class HeteroNodeView(object): |
| 16 | """A NodeView class to act as G.nodes for a DGLGraph.""" |
| 17 | |
| 18 | __slots__ = ["_graph", "_typeid_getter"] |
| 19 | |
| 20 | def __init__(self, graph, typeid_getter): |
| 21 | self._graph = graph |
| 22 | self._typeid_getter = typeid_getter |
| 23 | |
| 24 | def __getitem__(self, key): |
| 25 | if isinstance(key, slice): |
| 26 | # slice |
| 27 | if not ( |
| 28 | key.start is None and key.stop is None and key.step is None |
| 29 | ): |
| 30 | raise DGLError('Currently only full slice ":" is supported') |
| 31 | nodes = ALL |
| 32 | ntype = None |
| 33 | elif isinstance(key, tuple): |
| 34 | nodes, ntype = key |
| 35 | elif key is None or isinstance(key, str): |
| 36 | nodes = ALL |
| 37 | ntype = key |
| 38 | else: |
| 39 | nodes = key |
| 40 | ntype = None |
| 41 | ntid = self._typeid_getter(ntype) |
| 42 | return NodeSpace( |
| 43 | data=HeteroNodeDataView(self._graph, ntype, ntid, nodes) |
| 44 | ) |
| 45 | |
| 46 | def __call__(self, ntype=None): |
| 47 | """Return the nodes.""" |
| 48 | ntid = self._typeid_getter(ntype) |
| 49 | ret = F.arange( |
| 50 | 0, |
| 51 | self._graph._graph.num_nodes(ntid), |
| 52 | dtype=self._graph.idtype, |
| 53 | ctx=self._graph.device, |
| 54 | ) |
| 55 | return ret |
| 56 | |
| 57 | |
| 58 | class HeteroNodeDataView(MutableMapping): |