Decorate a whole library in Python

Question

I'm new to the ideas of decorators (and still trying to wrap my head around them), but I think I've come across a problem that would be well suited for them. I'd like to have class that is decorated across all of the functions in the math library. More specifically my class has two members, x and flag. When flag is true, I'd like the original math function to be called. When flag is false, I'd like to return None.

As a framework for what I'm asking here is the class:

import math

class num(object):
  def __init__(self, x, flag):
    self.x = x
    self.flag = flag

  def __float__(self):
    return float(self.x)

As a result, this works fine:

a = num(3, True)
print math.sqrt(a)

However this should (in my perfect world), return None:

b = num(4, False)
print math.sqrt(b)

Any suggestions or tips on how this would be possible to apply over a whole library of functions?

Answer 1

Here's the general idea...:

>>> class num(object):
...   def __init__(self, x, flag):
...     self.x = x
...     self.flag = flag
...   def __float__(self):
...     return float(self.x)
...   from functools import wraps
>>> def wrapper(f):
...   @wraps(f)
...   def wrapped(*a):
...     if not all(getattr(x, 'flag', True) for x in a):
...       return None
...     return f(*(getattr(x, 'x', x) for x in a))
...   return wrapped
... 
>>> import inspect
>>> import math
>>> for n, v in inspect.getmembers(math, inspect.isroutine):
...   setattr(math, n, wrapper(v))
... 

>>> a = num(3, True)
>>> print math.sqrt(a)
1.73205080757
>>> b = num(4, False)
>>> print math.sqrt(b)
None

Note that this wrapper also covers non-unary functions in math (returning None if any argument has a False .flag) and allows mixed calls thereof (with some args being instances of num and others being actual floats).

The key part, applicable to any "wrap all functions in a certain module" tasks, is using module inspect to get all the names and values of functions (built-in or not) in module math, and an explicit call to the wrapper (same semantics as the decorator syntax) to set that name to the wrapped value in the math module.