branching strategies are no longer variadic

Author: mrocklin

strategies/branch/core.py

@@ -35,16 +35,15 @@ def write(brl, x, result):
 
     return onaction(fn, write)
 
-def multiplex(*fns):
+@curry
+def multiplex(fns, x):
     """ Multiplex many branching rules into one """
-    def multiplex_brl(x):
-        seen = set([])
-        for brl in fns:
-            for nx in brl(x):
-                if nx not in seen:
-                    seen.add(nx)
-                    yield nx
-    return multiplex_brl
+    seen = set([])
+    for brl in fns:
+        for nx in brl(x):
+            if nx not in seen:
+                seen.add(nx)
+                yield nx
 
 @curry
 def condition(cond, fn, x):
@@ -70,33 +69,30 @@ def notempty(fn, x):
     if not yielded:
         yield x
 
-def do_one(*fns):
+@curry
+def do_one(fns, x):
     """ Execute one of the branching rules """
-    def do_one_brl(x):
-        yielded = False
-        for brl in fns:
-            for nx in brl(x):
-                yielded = True
-                yield nx
-            if yielded:
-                raise StopIteration()
-    return do_one_brl
+    yielded = False
+    for brl in fns:
+        for nx in brl(x):
+            yielded = True
+            yield nx
+        if yielded:
+            raise StopIteration()
 
-def chain(*fns):
+@curry
+def chain(fns, x):
     """
     Compose a sequence of fns so that they apply to the expr sequentially
     """
-    def chain_brl(x):
-        if not fns:
-            yield x
-            raise StopIteration()
-
-        head, tail = fns[0], fns[1:]
-        for nx in head(x):
-            for nnx in chain(*tail)(nx):
-                yield nnx
+    if not fns:
+        yield x
+        raise StopIteration()
 
-    return chain_brl
+    head, tail = fns[0], fns[1:]
+    for nx in head(x):
+        for nnx in chain(tail)(nx):
+            yield nnx
 
 
 @curry

strategies/branch/tests/test_core.py

@@ -51,7 +51,7 @@ def test_debug():
     assert '4' in log
 
 def test_multiplex():
-    brl = multiplex(posdec, branch5)
+    brl = multiplex([posdec, branch5])
     assert set(brl(3)) == set([2])
     assert set(brl(7)) == set([6, 8])
     assert set(brl(5)) == set([4, 6])
@@ -76,11 +76,11 @@ def ident_if_even(x):
     assert set(brl(5)) == set([5])
 
 def test_chain():
-    assert list(chain()(2)) == [2]  # identity
-    assert list(chain(inc, inc)(2)) == [4]
-    assert list(chain(branch5, inc)(4)) == [4]
-    assert set(chain(branch5, inc)(5)) == set([5, 7])
-    assert list(chain(inc, branch5)(5)) == [7]
+    assert list(chain([], 2)) == [2]  # identity
+    assert list(chain([inc, inc])(2)) == [4]
+    assert list(chain([branch5, inc])(4)) == [4]
+    assert set(chain([branch5, inc])(5)) == set([5, 7])
+    assert list(chain([inc, branch5])(5)) == [7]
 
 def test_onaction():
     L = []
@@ -103,6 +103,6 @@ def bad(expr):
         raise ValueError()
         yield False
 
-    assert list(do_one(inc)(3)) == [4]
-    assert list(do_one(inc, bad)(3)) == [4]
-    assert list(do_one(inc, posdec)(3)) == [4]
+    assert list(do_one([inc])(3)) == [4]
+    assert list(do_one([inc, bad])(3)) == [4]
+    assert list(do_one([inc, posdec])(3)) == [4]