In Python "everything is an object". Including classes, functions, and methods
on objects.
As such we can take any class, loop over all functions in that class and modify
them as needed.
Depending on the real code, the problem in the question might be better tackled
using decorators or meta-classes, depending on the dependencies of the wrapper
(what values does it need access to). I will not go into meta-classes as most needs for meta-classes can also be implemented using class-decorators, which are less error-prone.
As you mention in one of your comments that you may have several different classes that need to be wrapped the class-decorator solution might be a good candidate. This way you won't lose the inheritance tree of the wrapped class.
Here is an example not using either, but doing exactly as asked ;)
Using __new__
from functools import update_wrapper
class A:
def __init__(self):
self.a = 1
def forward(self, x):
"""
docstring (to demonstrate `update_wrapper`
"""
print("original forward")
return self.a + x
def main(self, x):
return self.forward(x) + 100
class Wrapper:
# Using __new__ instead of __init__ gives us complete control *how* the
# "Wrapper" instance is created. We use it to "pull in" methods from *A*
# and dynamically attach them to the `Wrapper` instance using `setattr`.
#
# Using __new__ is error-prone however, and using either meta-classes or
# even easier, decorators would be more maintainable.
def __new__(cls, A):
# instance will be our instance of thie `Wrapper` class. We start off
# with no defined functions, we will add those soon...
instance = super().__new__(cls)
instance.module = A
# We now walk over every "name" in the wrapped class
for funcname in dir(A):
# We skip anything starting with two underscores. They are most
# likely magic methods that we don't want to wrap with the
# additional code. The conditions what exactly we want to wrap, can
# be adapted as needed.
if funcname.startswith("__"):
continue
# We now need to get a reference to that attribute and check if
# it's callable. If not it is a member variable or something else
# and we can/should skip it.
func = getattr(A, funcname)
if not callable(func):
continue
# Now we "wrap" the function with our additional code. This is done
# in a separate function to keep __new__ somewhat clean
wrapped = Wrapper._wrap(func)
# After wrapping the function we can attach that new function ont
# our `Wrapper` instance
setattr(instance, funcname, wrapped)
return instance
@staticmethod
def _wrap(func):
"""
Wraps *func* with additional code.
"""
# we define a wrapper function. This will execute all additional code
# before and after the "real" function.
def wrapped(*args, **kwargs):
print("before-call:", func, args, kwargs)
output = func(*args, **kwargs)
print("after-call:", func, args, kwargs, output)
return output
# Use "update_wrapper" to keep docstrings and other function metadata
# intact
update_wrapper(wrapped, func)
# We can now return the wrapped function
return wrapped
class Demo2:
def foo(self):
print("yoinks")
classA = A()
otherInstance = Demo2()
wrapperA = Wrapper(classA)
wrapperB = Wrapper(otherInstance)
print(wrapperA.forward(10))
print(wrapperB.foo())
print("docstring is maintained: %r" % wrapperA.forward.__doc__)
Using a class decorator
With a class decorator, there is no need to override __new__ which can lead to hard to debug issues if not 100% properly implemented.
However, it has a key difference: It modifies the existing class "in-place", so the original class is lost in a way. Although you could keep a reference to it in the unlikely case that you need to.
Modifying this in-place does however also mean that you don't need to replace all your usages in your application with the new "wrapper" class, making it a lot easier to implement in an existing code-base and eliminating the risk that you forget to apply the wrapper on new instances.
from functools import update_wrapper
def _wrap(func):
"""
Wraps *func* with additional code.
"""
# we define a wrapper function. This will execute all additional code
# before and after the "real" function.
def wrapped(*args, **kwargs):
print("before-call:", func, args, kwargs)
output = func(*args, **kwargs)
print("after-call:", func, args, kwargs, output)
return output
# Use "update_wrapper" to keep docstrings and other function metadata
# intact
update_wrapper(wrapped, func)
# We can now return the wrapped function
return wrapped
def wrapper(cls):
for funcname in dir(cls):
# We skip anything starting with two underscores. They are most
# likely magic methods that we don't want to wrap with the
# additional code. The conditions what exactly we want to wrap, can
# be adapted as needed.
if funcname.startswith("__"):
continue
# We now need to get a reference to that attribute and check if
# it's callable. If not it is a member variable or something else
# and we can/should skip it.
func = getattr(cls, funcname)
if not callable(func):
continue
# Now we "wrap" the function with our additional code. This is done
# in a separate function to keep __new__ somewhat clean
wrapped = _wrap(func)
# After wrapping the function we can attach that new function ont
# our `Wrapper` instance
setattr(cls, funcname, wrapped)
return cls
@wrapper
class A:
def __init__(self):
self.a = 1
def forward(self, x):
"""
docstring (to demonstrate `update_wrapper`
"""
print("original forward")
return self.a + x
def main(self, x):
return self.forward(x) + 100
@wrapper
class Demo2:
def foo(self):
print("yoinks")
classA = A()
otherInstance = Demo2()
print(classA.forward(10))
print(otherInstance.foo())
print("docstring is maintained: %r" % classA.forward.__doc__)
Using function decorators
Another alternative, which diverges largely from the original question but may still prove insightful is using individual functions wrappers.
The code still used the same wrapper function, but here functions/methods are annotated individually.
This might give more flexibility by offering the possibility to leave some methods "unwrapped", but could easily lead to the wrapping code being executed more often than anticipated as demonstrated in the main() method.
from functools import update_wrapper
def wrap(func):
"""
Wraps *func* with additional code.
"""
# we define a wrapper function. This will execute all additional code
# before and after the "real" function.
def wrapped(*args, **kwargs):
print("before-call:", func, args, kwargs)
output = func(*args, **kwargs)
print("after-call:", func, args, kwargs, output)
return output
# Use "update_wrapper" to keep docstrings and other function metadata
# intact
update_wrapper(wrapped, func)
# We can now return the wrapped function
return wrapped
class A:
def __init__(self):
self.a = 1
@wrap
def forward(self, x):
"""
docstring (to demonstrate `update_wrapper`
"""
print("original forward")
return self.a + x
@wrap # careful: will be wrapped twice!
def main(self, x):
return self.forward(x) + 100
def foo(self):
print("yoinks")
classA = A()
print(">>> forward")
print(classA.forward(10))
print("<<< forward")
print(">>> main")
print(classA.main(100))
print("<<< main")
print(">>> foo")
print(classA.foo())
print("<<< foo")
