I have a number of calls I need to make to functions imported from a package (the package is a shared object file). However I need to do some pre-process / post - process step every time I make a call to a function from this package. Something like this:

import xyz

prepare()
xyz.foo(<args>)
done()

prepare()
xyz.bar(<args>)
done()

prepare()
xyz.foobar()
done()

Is there some way I can ask python to always invoke prepare() before I call a function from xyz module. And also to invoke done() after the call is complete? Writing prepare and done through my whole code seems redundant and messy. Appreciate your help!

CodePudding user response：

This is typically done with a context manager.

import contextlib

@contextlib.contextmanager
def with_preparation():
    prepare()
    yield
    done()

with preparation():
    xyz.foo(<args>)

with preparation():
    xyz.bar(<args>)

with preparation():
    xyz.foobar()

preparation defines a function that returns a context manager. The with statement works by invoking the context manager's __enter__ method, then executing the body, then ensuring that the context manager's __exit__ method is invoked before moving on (whether due to an exception being raised or the body completing normally).

contextlib.contextmanager provides a simple way to define a context manager using a generator function, rather than making you define a class with explicit __enter__ and __exit__ methods.

You mentioned you need this for every function in a particular module. Without exact details about the module, this may not be entirely correct, but you might be able to build up on it.

class XYZWrapper:
    def __getattr__(self, name):
        # Intentionally let an AttributeError propagate upwards
        f = getattr(xyz, name)
        def _(self, *args, **kwargs):
            prepare()
            return f(*args, **kwargs)
            done()
        setattr(XYZWrapper, name, _)
        return _

prepared = XYZWrapper()

prepared.foo(<args>)
prepared.bar(<args>)
prepared.foobar()
        
        

In short, any attribute access on the XYZWrapper instance tries to find an identical attribute on the xyz module, and if successful, defines a wrapper that calls prepare() and done() as needed and patches the XYZWrapper instance with the new wrapper.

CodePudding user response：

To extend @chepner's excellent answer, you can define your own class and use its __getattr__ function to create a function that wraps the actual module's function with your pre- and post-processing functions:

import typing
import xyz

def XYZWrapper():
    def __init__(self, pre, post):
        self.pre = pre
        self.post = post

    def __getattr__(self, a):
        func = getattr(xyz, a)
        if isinstance(func, typing.Callable):
            def wrapper(*args, **kwargs):
                self.pre()
                func(*args, **kwargs)
                self.post()
            return wrapper
        raise TypeError(f"'{type(func)}' object is not callable")

To use it, do

xyz = XYZWrapper(prepare, done)
xyz.foo(<args>)
...

Note that if you want to overwrite the xyz variable, you will need to put the wrapper class in a separate file.