1#
2# Copyright (c) 2017 Brian J. Kidney
3# All rights reserved.
4#
5# This software was developed by BAE Systems, the University of Cambridge
6# Computer Laboratory, and Memorial University under DARPA/AFRL contract
7# FA8650-15-C-7558 ("CADETS"), as part of the DARPA Transparent Computing
8# (TC) research program.
9#
10# Licensed under the Apache License, Version 2.0 (the "License");
11# you may not use this file except in compliance with the License.
12# You may obtain a copy of the License at
13#
14# http://www.apache.org/licenses/LICENSE-2.0
15#
16# Unless required by applicable law or agreed to in writing, software
17# distributed under the License is distributed on an "AS IS" BASIS,
18# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19# See the License for the specific language governing permissions and
20# limitations under the License.
21#
22
23
24class _Direction:
25    Successors, Predecessors = range(2)
26
27
28def get_roots(graph):
29    return [n for n, d in graph.in_degree().items() if d == 0]
30
31
32def get_leaves(graph):
33    return [n for n, d in graph.out_degree().items() if d == 0]
34
35
36def successors_by_generations(graph, node, generations):
37    return _neighbours_by_generations(graph, node, generations,
38                                      _Direction.Successors)
39
40
41def predecessors_by_generations(graph, node, generations):
42    return _neighbours_by_generations(graph, node, generations,
43                                      _Direction.Predecessors)
44
45
46def _neighbours_by_generations(graph, node, generations, direction):
47    neighbours = []
48    current_gen = []
49    current_gen.append(node)
50
51    for x in range(0, generations):
52        next_gen = []
53        for y in current_gen:
54            s = []
55            if direction == _Direction.Successors:
56                s = graph.successors(y)
57            elif direction == _Direction.Predecessors:
58                s = graph.predecessors(y)
59            neighbours.extend(s)
60            next_gen.extend(s)
61        current_gen = next_gen
62
63    return neighbours
64