CODEWITHSUNDEEP
pythonpython-3.x
backstreetrover
backstreetrover

Reputation: 323

Is there a way to do some pre-process/post-process every time I call a function in python

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!

Upvotes: 2

Views: 862

Answers (2)

chepner
chepner

Reputation: 531410

This is typically done with a context manager.

import contextlib

@contextlib.contextmanager
def 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.

Upvotes: 4

pho
pho

Reputation: 25489

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.

Upvotes: 2

Related Questions