McCarthy's amb operator in Python

From what I can gather after some reading [1], [2], [3]; McCarthy's amb operator, would allow you to give the possible values for variables; then give functions that constrain the values; and finally extract values for those variables that satisfy the constraints just by calling amb without any arguments.

I decided to work backwards from how I thought it should be used to an implementation.

Given the problem of finding Pythagorean triples for integers <=10, i.e define three variables x, y, and z that can have values 1..10 subject to the constraint that x*x +y*y == z*z. What are the possible values of x, y and z that satisfy the constraint? For an amb-type solution adapted a little to python, being amb-like is all in the way that the problem is expressed. I would like to write something like:

print("\nSmall Pythagorean triples problem")
# Set rage of global variables
x = amb(range(1,11))
y = amb(range(1,11))
z = amb(range(1,11))

# Extract solutions to global variables and print
for _dummy in amb( lambda x, y, z: x*x + y*y == z*z ):
print x, y, z

I took a little liberty in the way I had seen other implementations work, as I had seen examples where the call of amb with an argument of a function just set things up and 'registered' the function, then subsequent calls of amb() without any arguments, would set the global variables used in the lambda function to successive values that satisfy the constraint. I reasoned that calls to amb after registering the constraint function should follow the Python iterator protocol ( with an __iter__ and a next method), so all the values, if any, could be accessed in a loop. I liked the idea of automatically setting global variables so decided to give that a go (but I was prepared to jettison this setting global variables if it proved long-winded).

The solution follows. To modify global names, I noted that functions have f.func_globals which is the globals environment in which the function was created. I get this from the lambda expression. I also cheat a little in requiring the parameter names of the constraint function to be the global names that were previously assigned value ranges.

import itertools as _itertools

class Amb(object):
 def __init__(self):
     self._names2values   = {}       # set of values for each global name
     self._func           = None     # Boolean constraint function
     self._valueiterator  = None     # itertools.product of names values
     self._funcargnames   = None     # Constraint parameter names

 def __call__(self, arg=None):
     if hasattr(arg, 'func_globals'):
         ##
         ## Called with a constraint function. 
         ##
         globls = arg.func_globals
         # Names used in constraint
         argv = arg.__code__.co_varnames[:arg.__code__.co_argcount]
         for name in argv:
             if name not in self._names2values:
                 assert name in globls, \
                        "Global name %s not found in function globals" % name
                 self._names2values[name] = globls[name]
         # Gather the range of values of all names used in the constraint
         valuesets = [self._names2values[name] for name in argv]
         self._valueiterator = _itertools.product(*valuesets)
         self._func = arg
         self._funcargnames = argv
         return self
     elif arg is not None:
         ##
         ## Assume called with an iterable set of values
         ##
         arg = frozenset(arg)
         return arg
     else:
         ##
         ## blank call tries to return next solution
         ##
         return self._nextinsearch()

 def _nextinsearch(self):
     arg = self._func
     globls = arg.func_globals
     argv = self._funcargnames
     found = False
     for values in self._valueiterator:
         if arg(*values):
             # Set globals.
             found = True
             for n, v in zip(argv, values):
                 globls[n] = v
             break
     if not found: raise StopIteration
     return values

 def __iter__(self):
     return self

 def next(self):
     return self()

if __name__ == '__main__':
 if True:
     amb = Amb()

     print("\nSmall Pythagorean triples problem:")
     x = amb(range(1,11))
     y = amb(range(1,11))
     z = amb(range(1,11))

     for _dummy in amb( lambda x, y, z: x*x + y*y == z*z ):
         print x, y, z


 if True:
     amb = Amb()

     print("\nRosetta Code Amb problem:")
     w1 = amb(["the", "that", "a"])
     w2 = amb(["frog", "elephant", "thing"])
     w3 = amb(["walked", "treaded", "grows"])
     w4 = amb(["slowly", "quickly"])

     for _dummy in amb( lambda w1, w2, w3, w4: \
                          w1[-1] == w2[0] and \
                          w2[-1] == w3[0] and \
                          w3[-1] == w4[0] ):
         print w1, w2, w3, w4

 if True:
     amb = Amb()

     print("\nAmb problem from "
         "http://www.randomhacks.net/articles/2005/10/11/amb-operator:")
     x = amb([1, 2, 3])
     y = amb([4, 5, 6])

     for _dummy in amb( lambda x, y: x * y != 8 ):
         print x, y

Program output looks like this:

Small Pythagorean triples problem:
3 4 5
4 3 5
6 8 10
8 6 10

Rosetta Code Amb problem:
that thing grows slowly

Amb problem from http://www.randomhacks.net/articles/2005/10/11/amb-operator:
1 4
1 5
1 6
2 5
2 6
3 4
3 5
3 6

So, what you end up with is a declarative style of programming that is good as an exercise, but the solving algorithm is just a brute-force search over all permutations of possible inputs.