I'm trying to create a class that calls another class attributes using __getattr__ in order to wrap the class calls.
from aiohttp import ClientSession
from contextlib import asynccontextmanager
class SessionThrottler:
def __init__(self, session: ClientSession,
time_period: int, max_tasks: int):
self._obj = session
self._bucket = AsyncLeakyBucket(max_tasks=max_tasks,
time_period=time_period)
def __getattr__(self, name):
#asynccontextmanager
async def _do(*args, **kwargs):
async with self._bucket:
res = await getattr(self._obj, name)(*args, **kwargs)
yield res
return _do
async def close(self):
await self._obj.close()
So then I can do:
async def fetch(session: ClientSession):
async with session.get('http://localhost:5051') as resp:
_ = resp
session = ClientSession()
session_throttled = SessionThrottler(session, 4, 2)
await asyncio.gather(
*[fetch(session_trottled)
for _ in range(10)]
)
This code works fine but how can I do so that session_throttled is inferred as a ClientSession instead of SessionThrottler (kinda like functools.wraps) ?
I depends on what you need with "is inferred as".
Making ThrotledSessions instances of ClientSessions
The natural way of doing that with classes is trough inheritance - if your SessionThrotler inherits from ClientSession, it would be naturally be a ClientSession as well.
The "small downside" is that then __getattr__ would not work as expected, since is only called for attributes not found in the instance - and Python would "see" the original methods from ClientSession in your ThrotledSession object and call those instead.
Of course, that would also require you to statically inherit your class, and you may want it to work dynamically. (by statically, I mean
having to write class SessionThrotler(ClientSession): - or at least, if there is a finite number of different Session classes you want to wrap, write for each a subclass inheriting from ThrotledClass as well:
class ThrotledClientSession(ThrotledSession, ClientSession):
...
If that is something that would work for you, then it is a matter of fixing attribute access by creating __getattribute__ instead of __getattr__. The difference between both is that __getattribte__ emcompasses all of the attribute lookup steps, and is called at the beggning of the lookup. Whereas __getattr__ is called as part of the normal lookup (inside the standard algorithm for __getattribute__) when all else fails.
class SessionThrottlerMixin:
def __init__(self, session: ClientSession,
time_period: int, max_tasks: int):
self._bucket = AsyncLeakyBucket(max_tasks=max_tasks,
time_period=time_period)
def __getattribute__(self, name):
attr = super().__getattribute__(name)
if not name.startswith("_") or not callable(attr):
return attr
#asynccontextmanager
async def _do(*args, **kwargs):
async with self._bucket:
res = await attr(*args, **kwargs)
yield res
return _do
class ThrotledClientSession(SessionThrottlerMixin, ClientSession):
pass
If you are getting your CLientSession instances from other code, and don't want, or can't, replace the base class with this one, you can do that on the desired instances, by assigning to the __class__ attribute:
it works if ClientSession is a normal Python class, not inheriting from special bases like Python built-ins, not using __slots__ and a few other restrictions - the instance is "converted" to a ThrotledClientSession midflight (but you have to do the inheritance thing): session.__class__ = ThrottledClientSession.
Class assignemnt in this way won't run the new class __init__. Since you need the _bucket to be created, you could have a class method that would create the bucket and make the replacement - so, in the version with __getattribute__ add something like:
class SessionThrottler:
...
#classmethod
def _wrap(cls, instance, time_period: int, max_tasks: int):
cls.__class__ = cls
instance._bucket = AsyncLeakyBucket(max_tasks=max_tasks,
time_period=time_period)
return instance
...
throtled_session = ThrotledClientSession._wrap(session, 4, 2)
If you have a lot of parent classes that you want to wrap this way, and you don't want to declare a Throttledversion of it, this could be made dynamically - but I would only go that way if it were the only way to go. Declaring some 10 stub Thotled versions, 3 lines each, would be preferable.
Virtual Subclassing
If you can change the code of your ClientSession classes (and others you want to wrap) this is the least obtrusive way -
Python have an obscure OOP feature called Virtual Subclassing in which a class can be registered as a subclass of another, without real inheritance. However, the class that is to be the "parent" have to have abc.ABCMeta as its metaclass - otherwise this is really unobtrusive.
Here is how it works:
In [13]: from abc import ABC
In [14]: class A(ABC):
...: pass
...:
In [15]: class B: # don't inherit
...: pass
In [16]: A.register(B)
Out[16]: __main__.B
In [17]: isinstance(B(), A)
Out[17]: True
So, in your original code, if you can make ClientSession to inherit from abc.ABC (without any other change at all) - and then do:
ClientSession.register(SessionThrottler)
and it would just work (if the "inferred as" you mean has to do with the object type).
Note that if ClientSession and others have a different metaclass, adding abc.ABC as one of its bases will fail with a metaclass conflict. If you can change their code, this is still the better way to go: just create a collaborative metaclass that inherits from both metaclasses, and you are all set:
class Session(metaclass=IDontCare):
...
from abc import ABCMeta
class ColaborativeMeta(ABCMeta, Session.__class__):
pass
class ClientSession(Session, metaclass=ColaborativeMeta):
...
Type hinting
If you don't need "isinstance" to work, and just have to be the same for the typing system, then it is just a matter of using typing.cast:
import typing as T
...
session = ClientSession()
session_throttled = T.cast(ClientSession, SessionThrottler(session, 4, 2))
The object is untouched at runtime - just the same object, but from that point on, tools like mypy will consider it to be an instance of ClientSession.
Last, but not least, change the class name.
So, if by "inferred as" you don't mean that the wrapped class should be seen as an instance, but just care about the class name showing correctly in logs and such, you can just set the class __name__ attribute to whatever string you want:
class SessionThrottler:
...
SessionThrottelr.__name__ = ClientSession.__name__
Or just have an appropriate __repr__ method on the wrapper class:
class SessionThrottler:
...
def __repr__(self):
return repr(self._obj)
This solution is based on patching the object methods (instead of wrapping them) given to the context manager.
import asyncio
import functools
import contextlib
class Counter:
async def infinite(self):
cnt = 0
while True:
yield cnt
cnt += 1
await asyncio.sleep(1)
def limited_infinite(f, limit):
#functools.wraps(f)
async def inner(*a, **kw):
cnt = 0
async for res in f(*a, **kw):
yield res
if cnt == limit:
break
cnt += 1
return inner
#contextlib.contextmanager
def throttler(limit, counter):
orig = counter.infinite
counter.infinite = limited_infinite(counter.infinite, limit)
yield counter
counter.infinite = orig
async def main():
with throttler(5, Counter()) as counter:
async for x in counter.infinite():
print('res: ', x)
if __name__ == "__main__":
asyncio.run(main())
For your case the it means patching every relevant methods of ClientSession(http methods only probably). Not sure if it is better though.
Related
Let's say I have this
class Test():
def __init__(self, number):
self.number = number
await self.TestPrint()
async def TestPrint(self):
print(self.number)
As you can see this won't work since __init__ is not async and I cannot call await for the function
I want to be able to run TestPrint within __init__ assuming I want to maintain this function async
I also want this to have nothing to do with anything else other than the class (other function, other classes, main, etc.)
Thank you for your time.
Like chepner mentioned in the comments:
An asynchronous class method that creates the object and then calls the TestPrint method before returning it sounds more appropriate.
This is the preferred way above all else and why a lot of there are a lot of functions that initialize internal asyncio classes rather than you directly instantiating them.
That said, if you wish to have this close to the class you can use a #classmethod which can be async. Your code would look like this:
class Test():
def __init__(self, number):
self.number = number
async def TestPrint(self):
print(self.number)
#classmethod
async def with_print(cls, number):
self = cls(number)
await self.TestPrint()
return self
async def main():
t = await Test.with_print(123)
# 't' is now your Test instance.
...
I'm writing some testing utility library and may want to do the following
def assertSomething(test_case, ...):
test_case.assertEqual()
Is there a way to skip passing test_case around? e.g.
def assertSomething(...):
get_current_unittest_case().assertEqual()
AssertionError
If you just want to do some checks and fail with custom message,
raising AssertionError (via raise or assert) is the way to
go. By default, TestCase.failureException
is AssertionError, and
fail
(which internally is used by convenience methods of unittest.TestCase)
just raises it.
test_things.py:
import unittest
def check_is_zero(number):
assert number == 0, "{!r} is not 0".format(number)
def check_is_one(number):
if number != 1:
raise AssertionError("{!r} is not 1".format(number))
class NumbersTest(unittest.TestCase):
def test_one(self):
check_is_one(1)
def test_zero(self):
check_is_zero(0)
TestCase mixin
An easy and relatively readable way to add new assertions is to
make a “mixin class” that test cases will subclass. Sometimes it
is good enough.
testutils.py that contains mixin:
class ArithmeticsMixin:
def check_is_one(self, number):
self.assertEqual(number, 1)
test_thing.py, actual tests:
import unittest
import testutils
class NumbersTest(unittest.TestCase, testutils.ArithmeticsMixin):
def test_one(self):
self.check_is_one(1)
If there will be many mixin classes, use of base one may help:
import unittest
import testutils
class BaseTestCase(unittest.TestCase, testutils.ArithmeticsMixin):
"""Test case with additional methods for testing arithmetics."""
class NumbersTest(BaseTestCase):
def test_one(self):
self.check_is_one(1)
Thread local and TestCase subclass
Less readable is use of special base class,
thread local
(like global, but is aware of threads) and getter function.
testutils.py:
import unittest
import threading
_case_data = threading.local() # thread local
class ImproperlyUsed(Exception):
"""Raised if get_case is called without cooperation of BaseTestCase."""
def get_case(): # getter function
case = getattr(_case_data, "case", None)
if case is None:
raise ImproperlyUsed
return case
class BaseTestCase(unittest.TestCase): # base class
def run(self, *args, **kwargs):
_case_data.case = self
super().run(*args, **kwargs)
_case_data.case = None
def check_is_one(number):
case = get_case()
get_case().assertEqual(number, 1)
When test case is running, self (test case instance) is saved as
_case_data.case, so later inside check_is_one (or any other function
that is called from inside of test method and wants to access self)
get_case will be able to get reference to test case instance. Note
that after running _case_data.case is set to None, and if
get_case is called after that, ImproperlyUsed will be raised.
test_thing.py:
import testutils
def check_is_zero(number): # example of out-of-testutils function
testutils.get_case().assertEqual(number, 0)
class NumbersTest(testutils.BaseTestCase):
def test_one(self):
testutils.check_is_one(1)
def test_zero(self):
check_is_zero(0)
s̻̙ỵ̞̜͉͎̩s.̠͚̱̹̦̩͓_̢̬ge̡̯̳̖t̞͚̖̳f̜̪̩̪r̙̖͚̖̼͉̹a͏ṃ̡̹e̞̝̱̠̙
Finally, sys._getframe.
Let’s just h́o̞̓̐p̆̇̊̇e you don’t need it, because it is part of CPython,
not Python.
testutils.py:
import sys
import unittest
class ImproperlyUsed(Exception):
"""Raised if get_case can't get "self" TestCase instance."""
def get_case():
case = sys._getframe().f_back.f_back.f_locals.get("self")
if not isinstance(case, unittest.TestCase):
raise ImproperlyUsed
return case
def check_is_one(number):
case = get_case()
get_case().assertEqual(number, 1)
sys._getframe returns frame from the top of call stack, then few
frames below f_locals is taken, and there is self: instance of
unittest.TestCase. Like in previous implementation option, there is
sanity check, but here it is done with isinstance.
test_things.py:
import unittest
import testutils
def check_is_zero(number): # example of out-of-testutils function
testutils.get_case().assertEqual(number, 0)
class NumbersTest(unittest.TestCase):
def test_one(self):
testutils.check_is_one(1)
def test_zero(self):
check_is_zero(0)
If you just want to provide assertEqual for some new type,
take a look at addTypeEqualityFunc.
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)
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.
I'm writing a class that interfaces to a MoinMoin wiki via xmlrpc (simplified code follows):
class MoinMoin(object):
token = None
def __init__(self, url, username=None, password=None):
self.wiki = xmlrpclib.ServerProxy(url + '/?action=xmlrpc2')
if username and password:
self.token = self.wiki.getAuthToken(username, password)
# some sample methods:
def searchPages(self, regexp):
def getPage(self, page):
def putPage(self, page):
now each of my methods needs to call the relevant xmlrpc method alone if there isn't authentication involved, or to wrap it in a multicall if there's auth. Example:
def getPage(self, page):
if not self.token:
result = self.wiki.getPage(page)
else:
mc = xmlrpclib.MultiCall(self.wiki) # build an XML-RPC multicall
mc.applyAuthToken(self.token) # call 1
mc.getPage(page) # call 2
result = mc()[-1] # run both, keep result of the latter
return result
is there any nicer way to do it other than repeating that stuff for each and every method?
Since I have to call arbitrary methods, wrap them with stuff, then call the identically named method on another class, select relevant results and give them back, I suspect the solution would involve meta-classes or similar esoteric (for me) stuff. I should probably look at xmlrpclib sources and see how it's done, then maybe subclass their MultiCall to add my stuff...
But maybe I'm missing something easier. The best I've come out with is something like:
def _getMultiCall(self):
mc = xmlrpclib.MultiCall(self.wiki)
if self.token:
mc.applyAuthToken(self.token)
return mc
def fooMethod(self, x):
mc = self._getMultiCall()
mc.fooMethod(x)
return mc()[-1]
but it still repeats the same three lines of code for each and every method I need to implement, just changing the called method name. Any better?
Python function are objects so they can be passed quite easily to other function.
def HandleAuthAndReturnResult(self, method, arg):
mc = xmlrpclib.MultiCall(self.wiki)
if self.token:
mc.applyAuthToken(self.token)
method(mc, arg)
return mc()[-1]
def fooMethod(self, x):
HandleAuthAndReturnResult(xmlrpclib.MultiCall.fooMethod, x)
There may be other way but I think it should work. Of course, the arg part needs to be aligned with what is needed for the method but all your methods take one argument.
Edit: I didn't understand that MultiCall was a proxy object. Even if the real method call ultimately is the one in your ServerProxy, you should not pass this method object in case MultiCall ever overrides(define) it. In this case, you could use the getattribute method with the method name you want to call and then call the returned function object. Take care to handle the AttributeError exception.
Methods would now look like:
def HandleAuthAndReturnResult(self, methodName, arg):
mc = xmlrpclib.MultiCall(self.wiki)
if self.token:
mc.applyAuthToken(self.token)
try:
methodToCall = getattr(mc, methodName)
except AttributeError:
return None
methodToCall(arg)
return mc()[-1]
def fooMethod(self, x):
HandleAuthAndReturnResult('fooMethod', x)