"""Wrapping dags"""
from meshed import DAG
[docs]
class DDag(DAG):
wrappers = ()
def _call(self, *args, **kwargs):
if not self.wrappers:
return super()._call(*args, **kwargs)
else:
decorator = Line(*self.wrappers)
decorated_dag_call = decorator(super()._call)
return decorated_dag_call(*args, **kwargs)
def test_ddag():
def f(a, b=2):
return a + b
def g(f, c=3):
return f * c
# d = DDag([f, g])
d = DDag([f, g])
d.dot_digraph()
assert d(1, 2, 3) == 9 # can call
from i2 import Sig
assert str(Sig(d)) == '(a, b=2, c=3)' # has correct signature
def dec(func):
def _dec(*args, **kwargs):
print(func.__name__, args, kwargs)
return func(*args, **kwargs)
return _dec
def rev(func):
def _rev(*args, **kwargs):
assert not kwargs, "Can't have keyword arguments with rev"
return func(*args[::-1])
return _rev
d.wrappers = (dec,)
assert d(1, 2, 3) == 9
# prints: _call (1, 2, 3) {}
assert d(1, 2, c=3) == 9
# prints: _call (1, 2) {'c': 3}
d.wrappers = (dec, rev)
assert d(1, 2, 3) == 5
# prints: _call (3, 2, 1) {}