Python equivalent of Mathematica's Sow/Reap

Question

Suppose in Mathematica I define the following function:

f[list_] := Map[Prime[Sow[#]] &, list];

which outputs a list of prime numbers, such that if the input list has n at position i, then the output list will contain the nth prime number at position i. For example,

In[2]:= f[{1, 3, 4}]

Out[2]= {2, 5, 7}

Now, if for some reason (debugging, etc...) I want to check what values are being fed into the Prime function. Because of the Sow command in the function, I can do

In[3] := Reap[f[{1, 3, 4}]]

Out[3] := {{2, 5, 7}, {{1, 3, 4}}}

For more details on Sow/Reap, see the Wolfram Documentation. My question is, is there a natural Python equivalent of Mathematica's Sow and Reap functionality? In particular, is there a way to do this kind of thing without explicitly returning extra things from the python function you want to do it to, writing a second python function that is almost the same but returns something extra, or using global variables?

Answer 1

I came up with two ways to implement a rudimentary version of something like this, each with its own limitations. Here's the first version:

farm = []

def sower(func):
    def wrapped(*args, **kw):
        farm.append([])
        return func(*args, **kw)
    return wrapped

def sow(val):
    farm[-1].append(val)
    return val

def reap(val):
    return val, farm.pop()

You can use it like this (based on one of the examples from the Mathematica doc page):

>>> @sower
... def someSum():
...     return sum(sow(x**2 + 1) if (x**2 + 1) % 2 == 0 else x**2 + 1 for x in xrange(1, 11))
>>> someSum()
395
>>> reap(someSum())
(395, [2, 10, 26, 50, 82])

This has a number of limitations:

1. Any function that wants to use sow has to be decorated with the sower decorator. This means you can't use sow inside inside inline expressions like list comprehensions the way the Mathematica examples do. You might be able to hack this by inspecting the call stack, but it could get ugly.
2. Any values that are sown but not reaped get stored in the "farm" forever, so the farm will get bigger and bigger over time.
3. It doesn't have the "tag" abilities shown in the docs, although that wouldn't be too hard to add.

Writing this made me think of a simpler implementation with slightly different tradeoffs:

farm = []

def sow(val):
    if farm:
        farm[-1].append(val)
    return val

def reap(expr):
    farm.append([])
    val = expr()
    return val, farm.pop()

This one you can use like this, which is somewhat more similar to the Mathematica version:

>>> reap(lambda: sum(sow(x**2 + 1) if (x**2 + 1) % 2 == 0 else x**2 + 1 for x in xrange(1, 11)))
(395, [2, 10, 26, 50, 82])

This one doesn't require the decorator, and it cleans up reaped values, but it takes a no-argument function as its argument, which requires you to wrap your sowing expression in a function (here done with lambda). Also, this means that all sown values in any function called by the reaped expression will be inserted into the same list, which could result in weird ordering; I can't tell from the Mathematica docs if that's what Mathematica does or what.

Answer 2

Unfortunately, as far as I know, there's no simple or idiomatic equivalent in Python of "sow" and "reap". However, you might be able to fake it using a combination of generators and decorators like so:

def sow(func):
    class wrapper(object):
        def __call__(self, *args, **kwargs):
            output = list(func(*args, **kwargs))
            return output[-1]

        def reap(self, *args, **kwargs):
            output = list(func(*args, **kwargs))
            final = output[-1]
            intermediate = output[0:-1]
            return [final, intermediate]

    return wrapper()

@sow    
def f(seq, mul):
    yield seq
    yield mul
    yield [a * mul for a in seq]

print f([1, 2, 3], 4)         # [4, 8, 12]
print f.reap([1, 2, 3], 4)    # [[4, 8, 12], [[1, 2, 3], 4]]

However, compared to Mathematica, this method has a few limitations. First, the function has to be rewritten so it uses yield instead of return, turning it into a generator. The last value to be yielded would then be the final output.

It also doesn't have the same "exception"-like property that the docs describe. The decorator @sow is simply returning a class which fakes looking like a function, and adds an extra parameter reap.

---

An alternate solution might be to try using macropy. Since it directly manipulates the AST and bytecode of Python, you may be able to hack together direct support for something more in line with what you're looking for. The tracing macro looks vaguely similar in intent to what you want.