Flo
Flo

Reputation: 1387

How to get a method called (decorator?) after every object method

This is a question similar to How to call a method implicitly after every method call? but for python

Say I have a crawler class with some attributes (e.g. self.db) with a crawl_1(self, *args, **kwargs) and another one save_to_db(self, *args, **kwargs) which saves the crawling results to a database (self.db).

I want somehow to have save_to_db run after every crawl_1, crawl_2, etc. call. I've tried making this as a "global" util decorator but I don't like the result since it involves passing around self as an argument.

Upvotes: 1

Views: 1985

Answers (3)

larsks
larsks

Reputation: 312028

If you want to implicitly run a method after all of your crawl_* methods, the simplest solution may be to set up a metaclass that will programatically wrap the methods for you. Start with this, a simple wrapper function:

import functools

def wrapit(func):
    @functools.wraps(func)
    def _(self, *args, **kwargs):
        func(self, *args, **kwargs)
        self.save_to_db()

    return _

That's a basic decorator that wraps func, calling self.save_to_db() after calling func. Now, we set up a metaclass that will programatically apply this to specific methods:

class Wrapper (type):
    def __new__(mcls, name, bases, nmspc):
        for attrname, attrval in nmspc.items():
            if callable(attrval) and attrname.startswith('crawl_'):
                nmspc[attrname] = wrapit(attrval)

        return super(Wrapper, mcls).__new__(mcls, name, bases, nmspc)

This will iterate over the methods in the wrapped class, looking for method names that start with crawl_ and wrapping them with our decorator function.

Finally, the wrapped class itself, which declares Wrapper as a metaclass:

class Wrapped (object):
    __metaclass__ = Wrapper

    def crawl_1(self):
        print 'this is crawl 1'

    def crawl_2(self):
        print 'this is crawl 2'

    def this_is_not_wrapped(self):
        print 'this is not wrapped'

    def save_to_db(self):
        print 'saving to database'

Given the above, we get the following behavior:

>>> W = Wrapped()
>>> W.crawl_1()
this is crawl 1
saving to database
>>> W.crawl_2()
this is crawl 2
saving to database
>>> W.this_is_not_wrapped()
this is not wrapped
>>> 

You can see the our save_to_database method is being called after each of crawl_1 and crawl_2 (but not after this_is_not_wrapped).

The above works in Python 2. In Python 3, replase this:

class Wrapped (object):
    __metaclass__ = Wrapper

With:

class Wrapped (object, metaclass=Wrapper):

Upvotes: 6

Byte Commander
Byte Commander

Reputation: 6756

A decorator in Python looks like this, it's a method taking a single method as argument and returning another wrapper method that shall be called instead of the decorated one. Usually the wrapper "wraps" the decorated method, i.e. calls it before/after performing some other actions.

Example:

# define a decorator method:
def save_db_decorator(fn):

    # The wrapper method which will get called instead of the decorated method:
    def wrapper(self, *args, **kwargs):
        fn(self, *args, **kwargs)           # call the decorated method
        MyTest.save_to_db(self, *args, **kwargs)   # call the additional method

    return wrapper  # return the wrapper method

Now learn how to use it:

class MyTest:

    # The additional method called by the decorator:

    def save_to_db(self, *args, **kwargs):
        print("Saver")


    # The decorated methods:

    @save_db_decorator
    def crawl_1(self, *args, **kwargs):
        print("Crawler 1")

    @save_db_decorator
    def crawl_2(self, *args, **kwargs):
        print("Crawler 2")


# Calling the decorated methods:

my_test = MyTest()
print("Starting Crawler 1")
my_test.crawl_1()
print("Starting Crawler 1")
my_test.crawl_2()

This would output the following:

Starting Crawler 1
Crawler 1
Saver
Starting Crawler 1
Crawler 2
Saver

See this code running on ideone.com

Upvotes: 0

delanne
delanne

Reputation: 440

Something like this:

from functools import wraps

def my_decorator(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        print 'Calling decorated function'
        res = f(*args, **kwargs)
        obj = args[0] if len(args) > 0 else None
        if obj and hasattr(obj, "bar"):
            obj.bar()

    return wrapper

class MyClass(object):
    @my_decorator
    def foo(self, *args, **kwargs):
        print "Calling foo"

    def bar(self, *args, **kwargs):
        print "Calling bar"

@my_decorator
def example():
    print 'Called example function'

example()

obj = MyClass()
obj.foo()

It will give you the following output:

Calling decorated function
Called example function
Calling decorated function
Calling foo
Calling bar

Upvotes: 0

Related Questions