CODEWITHSUNDEEP
pythonclassmethodswrapper
Joe Jasinski
Joe Jasinski

Reputation: 10791

Python Wrap Class Method

I'm trying to create an object with a run method that will be wrapped by a _wrap_run method. I'd like to be able to call the method and it's wrapper by simply typing instance.run() and I'd like to be able to subclass the object so I can override the run() method and have it still execute the wrapper.

More simply put, I want people to be able to subclass A and override run() but still have calls to the run() method execute the wrapper function.

I'm having some difficulty with the mechanics of this. Does anyone have any suggestions regarding this approach?

class A:

    def run(self):
        print "Run A"
        return True

    def _wrap_run(self):
        print "PRE"
        return_value = self.run()
        print "POST"
        return return_value

    run = property(_wrap_run)


a = A()
a.run()
"""
Should Print: 
PRE
Run A
POST
"""


class B(A):

    def run(self):
        print "Run B"
        return True

b = B()
b.run()
"""
Should Print: 
PRE
Run B
POST
"""

Upvotes: 16

Views: 31420

Answers (5)

Nicholas Barrow
Nicholas Barrow

Reputation: 736

Here a while later, but if the method is a normal method, you could just use a regular decorator. Take the example below, which is a simplified SSH connector.

class SSH:

    def __init__(self) -> None:
        self.connected = False

    def connect(self) -> None:
        self.connected = True

    def call_only_if_connected(self) -> None:
        print("foo")

In this example, we only want the method call_only_if_connected to be called, as the name implies, if the SSH object that it is instantiated to has self.connected=True.

We can define a normal wrapper outside the class, and since self, or the instance of the class, is always passed as the first arg (args[0]), we can just make out check there. Consider:

def assert_connected(f):
    @wraps(f)
    def decorator(*args, **kwargs):
        if not args[0].connected:
            raise RuntimeError("attempted to call a method that requires a connection without an established connection")
        return f(*args, **kwargs)
    return decorator

class SSH:

    def __init__(self) -> None:
        self.connected = False

    def connect(self) -> None:
        self.connected = True

    @assert_connected
    def call_only_if_connected(self) -> None:
        print("foo")

Note in the above example that an AttributeError will be thrown if the object does not have a connected property, so it is important to define connected=False as soon as possible (i.e., as soon as you instantiate the object in __init__).

You could handle the improper connection how you'd like; in my example above, I am throwing a RuntimeError.

Upvotes: 1

S.Lott
S.Lott

Reputation: 391854

What other folks do

class A:
   def do_run( self ):
       """Must be overridden."""
       raise NotImplementedError
   def run( self, *args, **kw ):
       """Must not be overridden.
       You were warned.
       """
       print "PRE"
       return_value = self.do_run(*args, **kw)
       print "POST"
       return return_value

class B(A):
    def do_run(self):
        print "Run B"
        return True

That's usually sufficient.

If you want to worry about someone "breaking" this, stop now. Don't waste time worrying.

It's Python. We're all adults here. All the malicious sociopaths will break all you code by copying it, changing it, and then breaking it. No matter what you do, they'll just copy your code and modify it to break the clever bits.

Everyone else will read your comment and stick by your rules. If they want to use your module/package/framework, they will cooperate.

Upvotes: 3

agf
agf

Reputation: 176780

Use a Metaclass.

class MetaClass(type):
    @staticmethod
    def wrap(run):
        """Return a wrapped instance method"""
        def outer(self):
            print "PRE",
            return_value = run(self)
            print "POST"
            return return_value
        return outer
    def __new__(cls, name, bases, attrs):
        """If the class has a 'run' method, wrap it"""
        if 'run' in attrs:
            attrs['run'] = cls.wrap(attrs['run'])
        return super(MetaClass, cls).__new__(cls, name, bases, attrs)

class MyClass(object):
    """Use MetaClass to make this class"""
    __metaclass__ = MetaClass
    def run(self): print 'RUN',

myinstance = MyClass()

# Prints PRE RUN POST
myinstance.run()

Now if other people subclass MyClass, they will still get their run() methods wrapped.

Upvotes: 27

Gabi Purcaru
Gabi Purcaru

Reputation: 31524

What you have there is basically a decorator, so why not go ahead and implement _wrap_run as a decorator and apply it when subclassing the function?

Upvotes: 0

Cat Plus Plus
Cat Plus Plus

Reputation: 129764

Easiest way: make run the wrapper, and a private method be the overrideable one.

class A(object):
    def run(self):
        print "PRE"
        return_value = self._inner_run()
        print "POST"
        return return_value

    def _inner_run(self):
        print "Run A"
        return True

class B(A):
    def _inner_run(self):
        print "Run B"
        return True

Upvotes: 5

Related Questions