diff --git a/normal.py b/normal.py index a5f8a94..b18e0ec 100644 --- a/normal.py +++ b/normal.py @@ -59,6 +59,15 @@ def __call__(self, /, *args, **kwargs): result = function(result) return result + def __add__(self, other): + """Compose using '+' operator.""" + if isinstance(other, compose): + return compose(*self.functions, *other.functions) + elif callable(other): + return compose(*self.functions, other) + else: + raise ValueError("Can only compose callables and 'compose' instances.") + @_recursive_repr_if_available def __repr__(self): """Represent the composed function as an unambiguous string.""" diff --git a/test.py b/test.py index f59bd19..dd4a6ad 100644 --- a/test.py +++ b/test.py @@ -11,6 +11,46 @@ def g(s): assert f_g.functions == (g, f) # functions exposed in execution order +def test_compose_add_compose(): + def f(s): + return s + 'f' + def g(s): + return s + 'g' + f_wrapped = compose(f) + g_wrapped = compose(g) + f_g_wrapped = f_wrapped + g_wrapped + + assert f_g_wrapped('') == 'gf' # Remember that it's f(g(...)) so g runs first + assert f_g_wrapped.functions == (g, f) # functions exposed in execution order + + +def test_compose_add_function(): + def f(s): + return s + 'f' + def g(s): + return s + 'g' + f_wrapped = compose(f) + f_g_composed = f_wrapped + g + assert f_g_composed('') == 'gf' # Remember that it's f(g(...)) so g runs first + assert f_g_composed.functions == (g, f) # functions exposed in execution order + + +def test_compose_add_neither_compose_nor_function(): + def f(s): + return s + 'f' + f_wrapped = compose(f) + g = 'string' + + raised = False + + try: + f_g_wrapped = f_wrapped + g + except ValueError as err: + raised = True + finally: + assert raised + + def test_inlining(): def f(_): return None