lcm/workflows/graphflow/flow/graph.py

   1 # Copyright 2018 ZTE Corporation.
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #         http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14
  15 import logging
  16 from collections import deque
  17 from collections import OrderedDict
  18
  19 logger = logging.getLogger(__name__)
  20
  21
  22 class Graph(object):
  23
  24     def __init__(self, graph_dict=None):
  25         self.graph = OrderedDict()
  26         if graph_dict:
  27             for node, dep_nodes in list(graph_dict.items()):
  28                 self.add_node(node, dep_nodes)
  29
  30     def add_node(self, node, dep_nodes):
  31         if node not in self.graph:
  32             self.graph[node] = set()
  33         if isinstance(dep_nodes, list):
  34             for dep_node in dep_nodes:
  35                 if dep_node not in self.graph:
  36                     self.graph[dep_node] = set()
  37                 if dep_node not in self.graph[node]:
  38                     self.graph[node].add(dep_node)
  39
  40     def get_pre_nodes(self, node):
  41         return [k for k in self.graph if node in self.graph[k]]
  42
  43     def topo_sort(self):
  44         degree = {}
  45         for node in self.graph:
  46             degree[node] = 0
  47         for node in self.graph:
  48             for dependent in self.graph[node]:
  49                 degree[dependent] += 1
  50         queue = deque()
  51         for node in degree:
  52             if degree[node] == 0:
  53                 queue.appendleft(node)
  54         sort_list = []
  55         while queue:
  56             node = queue.pop()
  57             sort_list.append(node)
  58             for dependent in self.graph[node]:
  59                 degree[dependent] -= 1
  60                 if degree[dependent] == 0:
  61                     queue.appendleft(dependent)
  62         if len(sort_list) == len(self.graph):
  63             return sort_list
  64         else:
  65             return None
  66
  67     def to_dict(self):
  68         dict = {}
  69         for node, dependents in list(self.graph.items()):
  70             dict[node] = []
  71             for dep in dependents:
  72                 dict[node].append(dep)
  73         return dict