Adds the ability to concatenate independently created cursors.
This commit is contained in:
@ -29,17 +29,26 @@ class GraphCursor:
|
|||||||
def __rshift__(self, other):
|
def __rshift__(self, other):
|
||||||
""" Self >> Other """
|
""" Self >> Other """
|
||||||
|
|
||||||
|
# Allow to concatenate cursors.
|
||||||
|
if isinstance(other, GraphCursor):
|
||||||
|
chain = self.graph.add_chain(_input=self.last, _output=other.first)
|
||||||
|
return GraphCursor(chain.graph, first=self.first, last=other.last)
|
||||||
|
|
||||||
|
# If we get a partial graph, or anything with a node list, use that.
|
||||||
nodes = other.nodes if hasattr(other, "nodes") else [other]
|
nodes = other.nodes if hasattr(other, "nodes") else [other]
|
||||||
|
|
||||||
|
# Sometimes, we use ellipsis to show "pseudo-code". This is ok, but can't be executed.
|
||||||
if ... in nodes:
|
if ... in nodes:
|
||||||
raise NotImplementedError(
|
raise NotImplementedError(
|
||||||
"Expected something looking like a node, but got an Ellipsis (...). Did you forget to complete the graph?"
|
"Expected something looking like a node, but got an Ellipsis (...). Did you forget to complete the graph?"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# If there are nodes to add, create a new cursor after the chain is added to the graph.
|
||||||
if len(nodes):
|
if len(nodes):
|
||||||
chain = self.graph.add_chain(*nodes, _input=self.last, use_existing_nodes=True)
|
chain = self.graph.add_chain(*nodes, _input=self.last, use_existing_nodes=True)
|
||||||
return GraphCursor(chain.graph, first=self.first, last=chain.output)
|
return GraphCursor(chain.graph, first=self.first, last=chain.output)
|
||||||
|
|
||||||
|
# If we add nothing, then nothing changed.
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def __enter__(self):
|
def __enter__(self):
|
||||||
@ -49,7 +58,10 @@ class GraphCursor:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
def __eq__(self, other):
|
def __eq__(self, other):
|
||||||
return self.graph == other.graph and self.first == other.first and self.last == other.last
|
try:
|
||||||
|
return self.graph == other.graph and self.first == other.first and self.last == other.last
|
||||||
|
except AttributeError:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
class PartialGraph:
|
class PartialGraph:
|
||||||
|
|||||||
@ -158,3 +158,15 @@ def test_using_same_cursor_many_times_for_fork():
|
|||||||
assert g.outputs_of(c) == set()
|
assert g.outputs_of(c) == set()
|
||||||
assert g.outputs_of(d) == set()
|
assert g.outputs_of(d) == set()
|
||||||
assert g.outputs_of(e) == set()
|
assert g.outputs_of(e) == set()
|
||||||
|
|
||||||
|
|
||||||
|
def test_concat_branches():
|
||||||
|
a, b, c, d = get_pseudo_nodes(4)
|
||||||
|
g = Graph()
|
||||||
|
|
||||||
|
c0 = g.orphan() >> a >> b
|
||||||
|
c1 = g >> c >> d
|
||||||
|
c2 = c1 >> c0
|
||||||
|
|
||||||
|
assert c2.first == BEGIN
|
||||||
|
assert c2.last == g.index_of(b)
|
||||||
|
|||||||
Reference in New Issue
Block a user