I would use a decorator:
class Proxy(object):
def __init__(self, object_a):
self._object_a = object_a
def decorateEnterExit(self, obj, f):
def inner(*args, **kwargs):
with obj as _:
return f(*args, **kwargs)
return inner
def __getattribute__(self, name):
obj = object.__getattribute__(self, '_object_a')
dee = object.__getattribute__(self, 'decorateEnterExit')
return dee(obj, getattr(obj, name))
That way the with will be only evaluated when the function is executed. The way you did it it would first evaluate the with (including exiting it) and then the function would be returned to be called. Now we return a function which will itself enter the with and call the function inside.
>>> Proxy(A('Ax')).hello()
Enter the function
hello Ax!
Exit the function
Note that you need to return self
in your __enter__
method of A
:
def __enter__(self):
print 'Enter the function'
return self