enh: add mg(), a gfunc mathifier

Author: Technologicat

README.md

@@ -1437,6 +1437,8 @@ For convenience, we introduce some special cases:
 
 *Added in v0.13.1:* ``primes`` and ``fibonacci``.
 
+*Added in v0.14.0:* ``mg``, a decorator to mathify a gfunc, so that it will ``m()`` the generator instances it makes. Combo with ``imemoize`` for great justice, e.g. ``a = mg(imemoize(s(1, 2, ...)))``.
+
 We provide a compact syntax to create lazy constant, arithmetic, geometric and power sequences: ``s(...)``. Numeric (``int``, ``float``, ``mpmath``) and symbolic (SymPy) formats are supported. We avoid accumulating roundoff error when used with floating-point formats.
 
 We also provide arithmetic operation support for iterables (termwise). To make any iterable infix math aware, use ``m(iterable)``. The arithmetic is lazy; it just plans computations, returning a new lazy mathematical sequence. To extract values, iterate over the result. (Note this implies that expressions consisting of thousands of operations will overflow Python's call stack. In practice this shouldn't be a problem.)

unpythonic/mathseq.py

@@ -18,7 +18,7 @@
 (currently, the Fibonacci numbers and the prime numbers).
 """
 
-__all__ = ["s", "m", "almosteq",
+__all__ = ["s", "m", "mg", "almosteq",
            "sadd", "ssub", "sabs", "spos", "sneg", "sinvert", "smul", "spow",
            "struediv", "sfloordiv", "smod", "sdivmod",
            "sround", "strunc", "sfloor", "sceil",
@@ -27,6 +27,7 @@
            "fibonacci", "primes"]
 
 from itertools import repeat, takewhile, count
+from functools import wraps
 from operator import add as primitive_add, mul as primitive_mul, \
                      pow as primitive_pow, mod as primitive_mod, \
                      floordiv as primitive_floordiv, truediv as primitive_truediv, \
@@ -507,6 +508,23 @@ def __ror__(self, other):
         return sor(other, self)
     # TODO: conversion (bool, complex, int, float) and comparison operators? Do we want those?
 
+def mg(gfunc):
+    """Decorator: make gfunc m() the returned generator instances.
+
+    Return a new gfunc, which passes all its arguments to the original ``gfunc``.
+
+    Example::
+
+        a = mg(imemoize(s(1, 2, ...)))
+        assert last(take(5, a())) == 5
+        assert last(take(5, a())) == 5
+        assert last(take(5, a() + a())) == 10
+    """
+    @wraps(gfunc)
+    def mathify(*args, **kwargs):
+        return m(gfunc(*args, **kwargs))
+    return mathify
+
 # The *settings mechanism is used by round and pow.
 # These are recursive to support iterables containing iterables (e.g. an iterable of math sequences).
 def _make_termwise_stream_unop(op, *settings):

unpythonic/test/test_mathseq.py

@@ -4,7 +4,7 @@
 from math import floor, log2
 from operator import mul
 
-from ..mathseq import s, m, sadd, smul, spow, cauchyprod, primes, fibonacci
+from ..mathseq import s, m, mg, sadd, smul, spow, cauchyprod, primes, fibonacci
 from ..it import take, last
 from ..fold import scanl
 from ..gmemo import imemoize
@@ -233,6 +233,11 @@ def ulp(x):  # Unit in the Last Place
     factorials = imemoize(scanl(mul, 1, s(1, 2, ...)))  # 0!, 1!, 2!, ...
     assert last(take(6, factorials())) == 120
 
+    a = mg(imemoize(s(1, 2, ...)))
+    assert last(take(5, a())) == 5
+    assert last(take(5, a())) == 5
+    assert last(take(5, a() + a())) == 10
+
     print("All tests PASSED")
 
 if __name__ == '__main__':