Develop Reference

Override a function' sub-function from a decorator?

Let's consider this piece of code where I would like to create bar dynamically with a decorator
def foo():
def bar():
print "I am bar from foo"
print bar()
def baz():
def bar():
print "I am bar from baz"
print bar()
I thought I could create bar from the outside with a decorator:
def bar2():
print "I am super bar from foo"
setattr(foo, 'bar', bar2)
But the result is not what I was expecting (I would like to get I am super bar from foo:
>>> foo()
I am bar from foo
Is it possible to override a sub-function on an existing function with a decorator?
The actual use case
I am writing a wrapper for a library and to avoid boilerplate code I would like to simplify my work.
Each library function has a prefix lib_ and returns an error code. I would like to add the prefix to the current function and treat the error code. This could be as simple as this:
def call():
fname = __libprefix__ + inspect.stack()[1][3]
return_code = getattr(__lib__, fname)(*args)
if return_code < 0: raise LibError(fname, return_code)
def foo():
call()
The problem is that call might act differently in certain cases. Some library functions do not return an error_code so it would be easier to write it like
this:
def foo():
call(check_status=True)
Or much better in my opinion (this is the point where I started thinking about decorators):
#LibFunc(check_status=True)
def foo():
call()
In this last example I should declare call inside foo as a sub-function created dynamically by the decorator itself.
The idea was to use something like this:
class LibFunc(object):
def __init__(self,**kwargs):
self.kwargs = kwargs
def __call__(self, original_func):
decorator_self = self
def wrappee( *args, **kwargs):
def call(*args):
fname = __libprefix__ + original_func.__name__
return_code = getattr(__lib__, fname)(*args)
if return_code < 0: raise LibError(fname, return_code)
print original_func
print call
# <<<< The part that does not work
setattr(original_func, 'call', call)
# <<<<
original_func(*args,**kwargs)
return wrappee
Initially I was tempted to call the call inside the decorator itself to minimize the writing:
#LibFunc():
foo(): pass
Unfortunately, this is not an option since other things should sometime be done before and after the call:
#LibFunc():
foo(a,b):
value = c_float()
call(a, pointer(value), b)
return value.value
Another option that I thought about was to use SWIG, but again this is not an option because I will need to rebuild the existing library with the SWIG wrapping functions.
And last but not least, I may get inspiration from SWIG typemaps and declare my wrapper as this:
#LibFunc(check_exit = true, map = ('<a', '>c_float', '<c_int(b)')):
foo(a,b): pass
This looks like the best solution to me, but this is another topic and another question...

Are you married to the idea of a decorator? Because if your goal is bunch of module-level functions each of which wraps somelib.lib_somefunctionname, I don't see why you need one.
Those module-level names don't have to be functions, they just have to be callable. They could be a bunch of class instances, as long as they have a __call__ method.
I used two different subclasses to determine how to treat the return value:
#!/usr/bin/env python3
import libtowrap # Replace with the real library name.
class Wrapper(object):
'''
Parent class for all wrapped functions in libtowrap.
'''
def __init__(self, name):
self.__name__ = str(name)
self.wrapped_name = 'lib_' + self.__name__
self.wrapped_func = getattr(libtowrap, self.wrapped_name)
self.__doc__ = self.wrapped_func.__doc__
return
class CheckedWrapper(Wrapper):
'''
Wraps functions in libtowrap that return an error code that must
be checked. Negative return values indicate an error, and will
raise a LibError. Successful calls return None.
'''
def __call__(self, *args, **kwargs):
error_code = self.wrapped_func(*args, **kwargs)
if error_code < 0:
raise LibError(self.__name__, error_code)
return
class UncheckedWrapper(Wrapper):
'''
Wraps functions in libtowrap that return a useful value, as
opposed to an error code.
'''
def __call__(self, *args, **kwargs):
return self.wrapped_func(*args, **kwargs)
strict = CheckedWrapper('strict')
negative_means_failure = CheckedWrapper('negative_means_failure')
whatever = UncheckedWrapper('whatever')
negative_is_ok = UncheckedWrapper('negative_is_ok')
Note that the wrapper "functions" are assigned while the module is being imported. They are in the top-level module namespace, and not hidden by any if __name__ == '__main__' test.
They will behave like functions for most purposes, but there will be minor differences. For example, I gave each instance a __name__ that matches the name they're assigned to, not the lib_-prefixed name used in libtowrap... but I copied the original __doc__, which might refer to a prefixed name like lib_some_other_function. Also, testing them with isinstance will probably surprise people.
For more about decorators, and for many more annoying little discrepancies like the ones I mentioned above, see Graham Dumpleton's half-hour lecture "Advanced Methods for Creating Decorators" (PyCon US 2014; slides). He is the author of the wrapt module (Python Package Index; Git Hub; Read the Docs), which corrects all(?) of the usual decorator inconsistencies. It might solve your problem entirely (except for the old lib_-style names showing up in __doc__).

Given a method name as string, how to get the method instance?

I've a class Client which has many methods:
class Client:
def compute(self, arg):
#code
#more methods
All the methods of this class runs synchronously. I want to run them asynchronously. There are too many ways to accomplish this. But I'm thinking along these lines:
AsyncClient = make_async(Client) #make all methods of Client async, etc!
client = AsyncClient() #create an instance of AsyncClient
client.async_compute(arg) #compute asynchronously
client.compute(arg) #synchronous method should still exist!
Alright, that looks too ambitious, and I feel it can be done.
So far I've written this:
def make_async(cls):
class async_cls(cls): #derive from the given class
def __getattr__(self, attr):
for i in dir(cls):
if ("async_" + i) == attr:
#THE PROBLEM IS HERE
#how to get the method with name <i>?
return cls.__getattr__(i) # DOES NOT WORK
return async_cls
As you see the comment in the code above, the problem is to get the method given its name as string. How to do that? Once I get the method, I would wrap it in async_caller method, etc — the rest of the work I hope I can do myself.

The function __getattr__ just works with class instance, not class. Use getattr(cls, method_name) instead, it will solve the problem.
getattr(cls, method_name)

Asserting an object passed into a mocked method in the python mock

Say I have a class called Client that creates an object of the Request class and passes it to a method of a Connection object:
class Client(object):
def __init__(self, connection):
self._conn = connection
def sendText(plaintext):
self._conn.send(Request(0, plaintext))
And I want to assert the object passed into the Connection.send method to check its properties. I start with creating a mocked Connection class:
conn = Mock()
client = Client(conn)
client.sendText('some message')
And then I want something like:
conn.send.assert_called_with(
(Request,
{'type': 0, 'text': 'some message'})
)
Where 'type' and 'text' are properties of Request. Is there a way to do this in python's mock? All I found in the documentation were simple data examples.
I could have done it with mock.patch decorator by replacing the original 'send' method with a method which asserts the object's fields:
def patchedSend(self, req):
assert req.Type == 0
with mock.patch.object(Connection, 'send', TestClient.patchedSend):
...
but in this case I would have to define a separete mocked function for every method check and I couldn't check (without additional coding) if the function has been called at all.

You can get the last arguments to the mock with
request, = conn.send.call_args
and then assert properties about that. If you want facilities to express more sophisticated assertions about things, you can install PyHamcrest.
Note: Don't use assert in unit tests. Use assertion methods like assertEqual or assertTrue. Assertion methods can't be accidentally turned off, and they can give more useful messages than assert statements.

Well, I think that the easiest and better way of doing it, in this specific case, is to make a function to create requests and then mock it.
For instance, it'd be something like it:
class Client(object):
def __init__(self, connection):
self._conn = connection
def _create_request(self, plain_text):
return Request(0, plain_text)
def send_text(self, plaintext):
self._conn.send(self._create_request(plain_text))
And then, in the test, you could mock _create_request to return some specific value and then assert that send_text was called with it.
You could also get the parameters by call_args, as suggested, but I think it looks better this way.

Preventing function (or decorator) from being nested

I've got some code in a decorator that I only want run once. Many other functions (utility and otherwise) will be called later down the line, and I want to ensure that other functions that may have this decorator aren't accidentally used way down in the nest of function calls.
I also want to be able to check, at any point, whether or not the current code has been wrapped in the decorator or not.
I've written this, but I just wanted to see if anyone else can think of a better/more elegant solution than checking for the (hopefully!) unique function name in the stack.
import inspect
def my_special_wrapper(fn):
def my_special_wrapper(*args, **kwargs):
""" Do some magic, only once! """
# Check we've not done this before
for frame in inspect.stack()[1:]: # get stack, ignoring current!
if frame[3] == 'my_special_wrapper':
raise StandardError('Special wrapper cannot be nested')
# Do magic then call fn
# ...
fn(*args, **kwargs)
return my_special_wrapper
def within_special_wrapper():
""" Helper to check that the function has been specially wrapped """
for frame in inspect.stack():
if frame[3] == 'my_special_wrapper':
return True
return False
#my_special_wrapper
def foo():
print within_special_wrapper()
bar()
print 'Success!'
#my_special_wrapper
def bar():
pass
foo()

Here is an example of using a global for this task - in what I believe is a relatively safe way:
from contextlib import contextmanager
from functools import wraps
_within_special_context = False
#contextmanager
def flag():
global _within_special_context
_within_special_context = True
try:
yield
finally:
_within_special_context = False
#I'd argue this would be best replaced by just checking the variable, but
#included for completeness.
def within_special_wrapper():
return _within_special_context
def my_special_wrapper(f):
#wraps(f)
def internal(*args, **kwargs):
if not _within_special_context:
with flag():
...
f(*args, **kwargs)
else:
raise Exception("No nested calls!")
return internal
#my_special_wrapper
def foo():
print(within_special_wrapper())
bar()
print('Success!')
#my_special_wrapper
def bar():
pass
foo()
Which results in:
True
Traceback (most recent call last):
File "/Users/gareth/Development/so/test.py", line 39, in <module>
foo()
File "/Users/gareth/Development/so/test.py", line 24, in internal
f(*args, **kwargs)
File "/Users/gareth/Development/so/test.py", line 32, in foo
bar()
File "/Users/gareth/Development/so/test.py", line 26, in internal
raise Exception("No nested calls!")
Exception: No nested calls!
Using a context manager ensures that the variable is unset. You could just use try/finally, but if you want to modify the behaviour for different situations, the context manager can be made to be flexible and reusable.

The obvious solution is to have special_wrapper set a global flag, and just skip its magic if the flag is set.
This is about the only good use of a global variable - to allow a single piece of code to store information that is only used within that code, but which needs to survive the life of execution in that code.
It doesn't need to be set in global scope. The function could set the flag on itself, for example, or on any object or class, as long as nothing else will touch it.
As noted by Lattyware in comments, you'll want to use either a try/except, or perhaps even better, a context manager to ensure the variable is unset.
Update: If you need the wrapped code to be able to check if it is wrapped, then provide a function which returns the value of the flag. You might want to wrap it all up with a class for neatness.
Update 2: I see you're doing this for transaction management. There are probably already libraries which do this. I strongly recommend that you at least look at their code.

While my solution technically works, it requires a manual reset of the decorator, but you could very well modify things such that the outermost function is instead a class (with the instances being the wrappers of the decorated functions passed to it in __init__), and have reset() being called in __exit__(), which would then allow you to use the with statement to create the decorator to be usable only once within the context. Also note that it requires Python 3 due to the nonlocal keyword, but that can easily be adapted to 2.7 with a dict in place of the flag variable.
def once_usable(decorator):
"Apply this decorator function to the decorator you want to be usable only once until it is reset."
def outer_wrapper():
flag = False
def inner_wrapper(*args, **kwargs):
nonlocal flag
if not flag:
flag = True
return decorator(*args, **kwargs)
else:
print("Decorator currently unusable.") # raising an Error also works
def decorator_reset():
nonlocal flag
flag = False
return (inner_wrapper, decorator_reset)
return outer_wrapper()
Testing:
>>> def a(aa):
return aa*2
>>> def b(bb):
def wrapper(*args, **kwargs):
print("Decorated.")
return bb(*args, **kwargs)
return wrapper
>>> dec, reset = once_usable(b)
>>> aa = dec(a)
>>> aa(22)
Decorated.
44
>>> aaa = dec(a)
Decorator currently unusable.
>>> reset()
>>> aaa = dec(a)
>>> aaa(11)
Decorated.
22

In python, is there a good idiom for using context managers in setup/teardown

I am finding that I am using plenty of context managers in Python. However, I have been testing a number of things using them, and I am often needing the following:
class MyTestCase(unittest.TestCase):
def testFirstThing(self):
with GetResource() as resource:
u = UnderTest(resource)
u.doStuff()
self.assertEqual(u.getSomething(), 'a value')
def testSecondThing(self):
with GetResource() as resource:
u = UnderTest(resource)
u.doOtherStuff()
self.assertEqual(u.getSomething(), 'a value')
When this gets to many tests, this is clearly going to get boring, so in the spirit of SPOT/DRY (single point of truth/dont repeat yourself), I'd want to refactor those bits into the test setUp() and tearDown() methods.
However, trying to do that has lead to this ugliness:
def setUp(self):
self._resource = GetSlot()
self._resource.__enter__()
def tearDown(self):
self._resource.__exit__(None, None, None)
There must be a better way to do this. Ideally, in the setUp()/tearDown() without repetitive bits for each test method (I can see how repeating a decorator on each method could do it).
Edit: Consider the undertest object to be internal, and the GetResource object to be a third party thing (which we aren't changing).
I've renamed GetSlot to GetResource here—this is more general than specific case—where context managers are the way which the object is intended to go into a locked state and out.

How about overriding unittest.TestCase.run() as illustrated below? This approach doesn't require calling any private methods or doing something to every method, which is what the questioner wanted.
from contextlib import contextmanager
import unittest
#contextmanager
def resource_manager():
yield 'foo'
class MyTest(unittest.TestCase):
def run(self, result=None):
with resource_manager() as resource:
self.resource = resource
super(MyTest, self).run(result)
def test(self):
self.assertEqual('foo', self.resource)
unittest.main()
This approach also allows passing the TestCase instance to the context manager, if you want to modify the TestCase instance there.

Manipulating context managers in situations where you don't want a with statement to clean things up if all your resource acquisitions succeed is one of the use cases that contextlib.ExitStack() is designed to handle.
For example (using addCleanup() rather than a custom tearDown() implementation):
def setUp(self):
with contextlib.ExitStack() as stack:
self._resource = stack.enter_context(GetResource())
self.addCleanup(stack.pop_all().close)
That's the most robust approach, since it correctly handles acquisition of multiple resources:
def setUp(self):
with contextlib.ExitStack() as stack:
self._resource1 = stack.enter_context(GetResource())
self._resource2 = stack.enter_context(GetOtherResource())
self.addCleanup(stack.pop_all().close)
Here, if GetOtherResource() fails, the first resource will be cleaned up immediately by the with statement, while if it succeeds, the pop_all() call will postpone the cleanup until the registered cleanup function runs.
If you know you're only ever going to have one resource to manage, you can skip the with statement:
def setUp(self):
stack = contextlib.ExitStack()
self._resource = stack.enter_context(GetResource())
self.addCleanup(stack.close)
However, that's a bit more error prone, since if you add more resources to the stack without first switching to the with statement based version, successfully allocated resources may not get cleaned up promptly if later resource acquisitions fail.
You can also write something comparable using a custom tearDown() implementation by saving a reference to the resource stack on the test case:
def setUp(self):
with contextlib.ExitStack() as stack:
self._resource1 = stack.enter_context(GetResource())
self._resource2 = stack.enter_context(GetOtherResource())
self._resource_stack = stack.pop_all()
def tearDown(self):
self._resource_stack.close()
Alternatively, you can also define a custom cleanup function that accesses the resource via a closure reference, avoiding the need to store any extra state on the test case purely for cleanup purposes:
def setUp(self):
with contextlib.ExitStack() as stack:
resource = stack.enter_context(GetResource())
def cleanup():
if necessary:
one_last_chance_to_use(resource)
stack.pop_all().close()
self.addCleanup(cleanup)

pytest fixtures are very close to your idea/style, and allow for exactly what you want:
import pytest
from code.to.test import foo
#pytest.fixture(...)
def resource():
with your_context_manager as r:
yield r
def test_foo(resource):
assert foo(resource).bar() == 42

The problem with calling __enter__ and __exit__ as you did, is not that you have done so: they can be called outside of a with statement. The problem is that your code has no provision to call the object's __exit__ method properly if an exception occurs.
So, the way to do it is to have a decorator that will wrap the call to your original method in a withstatement. A short metaclass can apply the decorator transparently to all methods named test* in the class -
# -*- coding: utf-8 -*-
from functools import wraps
import unittest
def setup_context(method):
# the 'wraps' decorator preserves the original function name
# otherwise unittest would not call it, as its name
# would not start with 'test'
#wraps(method)
def test_wrapper(self, *args, **kw):
with GetSlot() as slot:
self._slot = slot
result = method(self, *args, **kw)
delattr(self, "_slot")
return result
return test_wrapper
class MetaContext(type):
def __new__(mcs, name, bases, dct):
for key, value in dct.items():
if key.startswith("test"):
dct[key] = setup_context(value)
return type.__new__(mcs, name, bases, dct)
class GetSlot(object):
def __enter__(self):
return self
def __exit__(self, *args, **kw):
print "exiting object"
def doStuff(self):
print "doing stuff"
def doOtherStuff(self):
raise ValueError
def getSomething(self):
return "a value"
def UnderTest(*args):
return args[0]
class MyTestCase(unittest.TestCase):
__metaclass__ = MetaContext
def testFirstThing(self):
u = UnderTest(self._slot)
u.doStuff()
self.assertEqual(u.getSomething(), 'a value')
def testSecondThing(self):
u = UnderTest(self._slot)
u.doOtherStuff()
self.assertEqual(u.getSomething(), 'a value')
unittest.main()
(I also included mock implementations of "GetSlot" and the methods and functions in your example so that I myself could test the decorator and metaclass I am suggesting on this answer)

I'd argue you should separate your test of the context manager from your test of the Slot class. You could even use a mock object simulating the initialize/finalize interface of slot to test the context manager object, and then test your slot object separately.
from unittest import TestCase, main
class MockSlot(object):
initialized = False
ok_called = False
error_called = False
def initialize(self):
self.initialized = True
def finalize_ok(self):
self.ok_called = True
def finalize_error(self):
self.error_called = True
class GetSlot(object):
def __init__(self, slot_factory=MockSlot):
self.slot_factory = slot_factory
def __enter__(self):
s = self.s = self.slot_factory()
s.initialize()
return s
def __exit__(self, type, value, traceback):
if type is None:
self.s.finalize_ok()
else:
self.s.finalize_error()
class TestContextManager(TestCase):
def test_getslot_calls_initialize(self):
g = GetSlot()
with g as slot:
pass
self.assertTrue(g.s.initialized)
def test_getslot_calls_finalize_ok_if_operation_successful(self):
g = GetSlot()
with g as slot:
pass
self.assertTrue(g.s.ok_called)
def test_getslot_calls_finalize_error_if_operation_unsuccessful(self):
g = GetSlot()
try:
with g as slot:
raise ValueError
except:
pass
self.assertTrue(g.s.error_called)
if __name__ == "__main__":
main()
This makes code simpler, prevents concern mixing and allows you to reuse the context manager without having to code it in many places.