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

You can use decorators for this, although you won't need the @decorator syntax.

The following code imports each function you list from the math module into the current module's namespace, wrapping it in the defined decorator. It should give you the basic idea.

from functools import wraps
def check_flag(func):
    @wraps(func)
    def _exec(x, *args, **kw):
        if getattr(x, 'flag', False):
            return None

        return func(x, *args, **kw)

    return _exec

import sys, math
_module = sys.modules[__name__]
for func in ('exp', 'log', 'sqrt'):
    setattr(_module, func, check_flag(getattr(math, func)))

You could automate the listing of functions defined in the math module, as Alex demonstrates, but I think explicitly wrapping just the functions you're interested in using is a better way to go.

Answer 2

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.