|
|
from __future__ import unicode_literals |
|
|
|
|
|
from ._utils import get_hash, get_hash_int |
|
|
from builtins import object |
|
|
from collections import namedtuple |
|
|
|
|
|
|
|
|
class DagNode(object): |
|
|
"""Node in a directed-acyclic graph (DAG). |
|
|
|
|
|
Edges: |
|
|
DagNodes are connected by edges. An edge connects two nodes with a label for each side: |
|
|
- ``upstream_node``: upstream/parent node |
|
|
- ``upstream_label``: label on the outgoing side of the upstream node |
|
|
- ``downstream_node``: downstream/child node |
|
|
- ``downstream_label``: label on the incoming side of the downstream node |
|
|
|
|
|
For example, DagNode A may be connected to DagNode B with an edge labelled "foo" on A's side, and "bar" on B's |
|
|
side: |
|
|
|
|
|
_____ _____ |
|
|
| | | | |
|
|
| A >[foo]---[bar]> B | |
|
|
|_____| |_____| |
|
|
|
|
|
Edge labels may be integers or strings, and nodes cannot have more than one incoming edge with the same label. |
|
|
|
|
|
DagNodes may have any number of incoming edges and any number of outgoing edges. DagNodes keep track only of |
|
|
their incoming edges, but the entire graph structure can be inferred by looking at the furthest downstream |
|
|
nodes and working backwards. |
|
|
|
|
|
Hashing: |
|
|
DagNodes must be hashable, and two nodes are considered to be equivalent if they have the same hash value. |
|
|
|
|
|
Nodes are immutable, and the hash should remain constant as a result. If a node with new contents is required, |
|
|
create a new node and throw the old one away. |
|
|
|
|
|
String representation: |
|
|
In order for graph visualization tools to show useful information, nodes must be representable as strings. The |
|
|
``repr`` operator should provide a more or less "full" representation of the node, and the ``short_repr`` |
|
|
property should be a shortened, concise representation. |
|
|
|
|
|
Again, because nodes are immutable, the string representations should remain constant. |
|
|
""" |
|
|
|
|
|
def __hash__(self): |
|
|
"""Return an integer hash of the node.""" |
|
|
raise NotImplementedError() |
|
|
|
|
|
def __eq__(self, other): |
|
|
"""Compare two nodes; implementations should return True if (and only if) hashes match.""" |
|
|
raise NotImplementedError() |
|
|
|
|
|
def __repr__(self, other): |
|
|
"""Return a full string representation of the node.""" |
|
|
raise NotImplementedError() |
|
|
|
|
|
@property |
|
|
def short_repr(self): |
|
|
"""Return a partial/concise representation of the node.""" |
|
|
raise NotImplementedError() |
|
|
|
|
|
@property |
|
|
def incoming_edge_map(self): |
|
|
"""Provides information about all incoming edges that connect to this node. |
|
|
|
|
|
The edge map is a dictionary that maps an ``incoming_label`` to ``(outgoing_node, outgoing_label)``. Note that |
|
|
implicity, ``incoming_node`` is ``self``. See "Edges" section above. |
|
|
""" |
|
|
raise NotImplementedError() |
|
|
|
|
|
|
|
|
DagEdge = namedtuple( |
|
|
'DagEdge', |
|
|
[ |
|
|
'downstream_node', |
|
|
'downstream_label', |
|
|
'upstream_node', |
|
|
'upstream_label', |
|
|
'upstream_selector', |
|
|
], |
|
|
) |
|
|
|
|
|
|
|
|
def get_incoming_edges(downstream_node, incoming_edge_map): |
|
|
edges = [] |
|
|
for downstream_label, upstream_info in list(incoming_edge_map.items()): |
|
|
upstream_node, upstream_label, upstream_selector = upstream_info |
|
|
edges += [ |
|
|
DagEdge( |
|
|
downstream_node, |
|
|
downstream_label, |
|
|
upstream_node, |
|
|
upstream_label, |
|
|
upstream_selector, |
|
|
) |
|
|
] |
|
|
return edges |
|
|
|
|
|
|
|
|
def get_outgoing_edges(upstream_node, outgoing_edge_map): |
|
|
edges = [] |
|
|
for upstream_label, downstream_infos in sorted(outgoing_edge_map.items()): |
|
|
for downstream_info in downstream_infos: |
|
|
downstream_node, downstream_label, downstream_selector = downstream_info |
|
|
edges += [ |
|
|
DagEdge( |
|
|
downstream_node, |
|
|
downstream_label, |
|
|
upstream_node, |
|
|
upstream_label, |
|
|
downstream_selector, |
|
|
) |
|
|
] |
|
|
return edges |
|
|
|
|
|
|
|
|
class KwargReprNode(DagNode): |
|
|
"""A DagNode that can be represented as a set of args+kwargs. |
|
|
""" |
|
|
|
|
|
@property |
|
|
def __upstream_hashes(self): |
|
|
hashes = [] |
|
|
for downstream_label, upstream_info in list(self.incoming_edge_map.items()): |
|
|
upstream_node, upstream_label, upstream_selector = upstream_info |
|
|
hashes += [ |
|
|
hash(x) |
|
|
for x in [ |
|
|
downstream_label, |
|
|
upstream_node, |
|
|
upstream_label, |
|
|
upstream_selector, |
|
|
] |
|
|
] |
|
|
return hashes |
|
|
|
|
|
@property |
|
|
def __inner_hash(self): |
|
|
props = {'args': self.args, 'kwargs': self.kwargs} |
|
|
return get_hash(props) |
|
|
|
|
|
def __get_hash(self): |
|
|
hashes = self.__upstream_hashes + [self.__inner_hash] |
|
|
return get_hash_int(hashes) |
|
|
|
|
|
def __init__(self, incoming_edge_map, name, args, kwargs): |
|
|
self.__incoming_edge_map = incoming_edge_map |
|
|
self.name = name |
|
|
self.args = args |
|
|
self.kwargs = kwargs |
|
|
self.__hash = self.__get_hash() |
|
|
|
|
|
def __hash__(self): |
|
|
return self.__hash |
|
|
|
|
|
def __eq__(self, other): |
|
|
return hash(self) == hash(other) |
|
|
|
|
|
@property |
|
|
def short_hash(self): |
|
|
return '{:x}'.format(abs(hash(self)))[:12] |
|
|
|
|
|
def long_repr(self, include_hash=True): |
|
|
formatted_props = ['{!r}'.format(arg) for arg in self.args] |
|
|
formatted_props += [ |
|
|
'{}={!r}'.format(key, self.kwargs[key]) for key in sorted(self.kwargs) |
|
|
] |
|
|
out = '{}({})'.format(self.name, ', '.join(formatted_props)) |
|
|
if include_hash: |
|
|
out += ' <{}>'.format(self.short_hash) |
|
|
return out |
|
|
|
|
|
def __repr__(self): |
|
|
return self.long_repr() |
|
|
|
|
|
@property |
|
|
def incoming_edges(self): |
|
|
return get_incoming_edges(self, self.incoming_edge_map) |
|
|
|
|
|
@property |
|
|
def incoming_edge_map(self): |
|
|
return self.__incoming_edge_map |
|
|
|
|
|
@property |
|
|
def short_repr(self): |
|
|
return self.name |
|
|
|
|
|
|
|
|
def topo_sort(downstream_nodes): |
|
|
marked_nodes = [] |
|
|
sorted_nodes = [] |
|
|
outgoing_edge_maps = {} |
|
|
|
|
|
def visit( |
|
|
upstream_node, |
|
|
upstream_label, |
|
|
downstream_node, |
|
|
downstream_label, |
|
|
downstream_selector=None, |
|
|
): |
|
|
if upstream_node in marked_nodes: |
|
|
raise RuntimeError('Graph is not a DAG') |
|
|
|
|
|
if downstream_node is not None: |
|
|
outgoing_edge_map = outgoing_edge_maps.get(upstream_node, {}) |
|
|
outgoing_edge_infos = outgoing_edge_map.get(upstream_label, []) |
|
|
outgoing_edge_infos += [ |
|
|
(downstream_node, downstream_label, downstream_selector) |
|
|
] |
|
|
outgoing_edge_map[upstream_label] = outgoing_edge_infos |
|
|
outgoing_edge_maps[upstream_node] = outgoing_edge_map |
|
|
|
|
|
if upstream_node not in sorted_nodes: |
|
|
marked_nodes.append(upstream_node) |
|
|
for edge in upstream_node.incoming_edges: |
|
|
visit( |
|
|
edge.upstream_node, |
|
|
edge.upstream_label, |
|
|
edge.downstream_node, |
|
|
edge.downstream_label, |
|
|
edge.upstream_selector, |
|
|
) |
|
|
marked_nodes.remove(upstream_node) |
|
|
sorted_nodes.append(upstream_node) |
|
|
|
|
|
unmarked_nodes = [(node, None) for node in downstream_nodes] |
|
|
while unmarked_nodes: |
|
|
upstream_node, upstream_label = unmarked_nodes.pop() |
|
|
visit(upstream_node, upstream_label, None, None) |
|
|
return sorted_nodes, outgoing_edge_maps |
|
|
|
