Model.py 文件源码

def get_nodes(self, name, next_layer=False, last_layer=False, type=None):
        if type is None:
            name2node = {node.name: node for node in self.nodes()}
        else:
            name2node = {node.name: node for node in self.nodes() if node.type in type}
        assert name in name2node.keys(), " Name must be uniqiue"
        node = name2node[name]
        if next_layer:
            if type is None:
                return self.successors(node)
            else:
                poss_list, begin = [], False
                for poss in nx.topological_sort(self):
                    if poss == node:
                        begin = True
                        continue
                    if begin and poss in name2node.values():
                        poss_list.append(poss)
                return [poss_list[0]]
        elif last_layer:
            return self.predecessors(node)
        else:
            return [node]