Last active
August 24, 2024 12:08
-
-
Save anatoly-kussul/f2d7444443399e51e2f83a76f112364d to your computer and use it in GitHub Desktop.
Python sync-async decorator factory
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| class SyncAsyncDecoratorFactory: | |
| """ | |
| Factory creates decorator which can wrap either a coroutine or function. | |
| To return something from wrapper use self._return | |
| If you need to modify args or kwargs, you can yield them from wrapper | |
| """ | |
| def __new__(cls, *args, **kwargs): | |
| instance = super().__new__(cls) | |
| # This is for using decorator without parameters | |
| if len(args) == 1 and not kwargs and (inspect.iscoroutinefunction(args[0]) or inspect.isfunction(args[0])): | |
| instance.__init__() | |
| return instance(args[0]) | |
| return instance | |
| class ReturnValue(Exception): | |
| def __init__(self, return_value): | |
| self.return_value = return_value | |
| @contextmanager | |
| def wrapper(self, *args, **kwargs): | |
| raise NotImplementedError | |
| @classmethod | |
| def _return(cls, value): | |
| raise cls.ReturnValue(value) | |
| def __call__(self, func): | |
| @wraps(func) | |
| def call_sync(*args, **kwargs): | |
| try: | |
| with self.wrapper(*args, **kwargs) as new_args: | |
| if new_args: | |
| args, kwargs = new_args | |
| return self.func(*args, **kwargs) | |
| except self.ReturnValue as r: | |
| return r.return_value | |
| @wraps(func) | |
| async def call_async(*args, **kwargs): | |
| try: | |
| with self.wrapper(*args, **kwargs) as new_args: | |
| if new_args: | |
| args, kwargs = new_args | |
| return await self.func(*args, **kwargs) | |
| except self.ReturnValue as r: | |
| return r.return_value | |
| self.func = func | |
| return call_async if inspect.iscoroutinefunction(func) else call_sync |
I've been trying to use your code to avoir code duplication for sync and async functions decorators but couldn't get it to work.
I'm trying to implement a Time Bounbed LRU cache decorator (similar as functools.lru_cachebut with time expiring keys). Below my attempt:
from functools import wraps
import asyncio
from contextlib import contextmanager
import collections
import time
class SyncAsyncDecoratorFactory:
@contextmanager
def wrapper(self, *args, **kwargs):
raise NotImplementedError("Please call this method from subclasses")
def __call__(self, func):
@wraps(func)
def sync_wrapper(*args, **kwargs):
with self.wrapper(*args, **kwargs) as res:
if res is None:
return func(*args, **kwargs)
@wraps(func)
async def async_wrapper(*args, **kwargs):
with self.wrapper(*args, **kwargs) as res:
return await func(*args, **kwargs)
if asyncio.iscoroutinefunction(func):
return async_wrapper
else:
return sync_wrapper
class TimedLRU(SyncAsyncDecoratorFactory):
def __init__(self, max_size=128, max_age=30):
super().__init__()
self.max_size = max_size
self.max_age = max_age
self._cache = collections.OrderedDict()
self._sentinel = object()
@contextmanager
def wrapper(self, *args, **kwargs):
k = args + (self._sentinel,) + tuple(sorted(kwargs.items()))
if k in self._cache:
print("hit")
self._cache.move_to_end(k)
res, ts = self._cache[k]
if time.time() - ts <= self.max_age:
yield
return res
res = yield
self._cache[k] = (res, time.time())
if len(self._cache) > self.max_size:
self._cache.popitem(0)
return res
@TimedLRU()
def foobar(s):
print("in decorated function")
return "foo %s bar" % s
a = foobar("hey")
print(a)
b = foobar("hey")
print(b)Any help how to achieve this ? the _cacheattribute has Noneas a value and in decorated functionis printed twice instead of once.
Hey, sorry for late response.
The reason you are not getting what you want is that you can't directly return from context manager.
Here is slightly modified version that does what you want, not the best solution probably, just my first look at that, but you should get an idea:
from functools import wraps
import asyncio
from contextlib import contextmanager
import collections
import time
import inspect
class SyncAsyncDecoratorFactory:
"""
Factory creates decorator which can wrap either a coroutine or function.
To return something from wrapper use self._return
If you need to modify args or kwargs, you can yield them from wrapper
"""
def __new__(cls, *args, **kwargs):
instance = super().__new__(cls)
# This is for using decorator without parameters
if len(args) == 1 and not kwargs and (inspect.iscoroutinefunction(args[0]) or inspect.isfunction(args[0])):
instance.__init__()
return instance(args[0])
return instance
class ReturnValue(Exception):
def __init__(self, return_value):
self.return_value = return_value
@contextmanager
def wrapper(self, *args, **kwargs):
raise NotImplementedError
@classmethod
def _return(cls, value):
raise cls.ReturnValue(value)
def __call__(self, func):
@wraps(func)
def call_sync(*args, **kwargs):
try:
with self.wrapper(*args, **kwargs) as new_args:
if new_args:
args, kwargs = new_args
self._return(self.func(*args, **kwargs))
except self.ReturnValue as r:
return r.return_value
@wraps(func)
async def call_async(*args, **kwargs):
try:
with self.wrapper(*args, **kwargs) as new_args:
if new_args:
args, kwargs = new_args
self._return(await self.func(*args, **kwargs))
except self.ReturnValue as r:
return r.return_value
self.func = func
return call_async if inspect.iscoroutinefunction(func) else call_sync
class TimedLRU(SyncAsyncDecoratorFactory):
def __init__(self, max_size=128, max_age=30):
super().__init__()
self.max_size = max_size
self.max_age = max_age
self._cache = collections.OrderedDict()
self._sentinel = object()
@contextmanager
def wrapper(self, *args, **kwargs):
k = args + (self._sentinel,) + tuple(sorted(kwargs.items()))
if k in self._cache:
print("hit")
self._cache.move_to_end(k)
res, ts = self._cache[k]
if time.time() - ts <= self.max_age:
self._return(res)
try:
yield
except self.ReturnValue as r:
self._cache[k] = (r.return_value, time.time())
if len(self._cache) > self.max_size:
self._cache.popitem(0)
raise
@TimedLRU()
def foobar(s):
print("in decorated function")
return "foo %s bar" % s
a = foobar("hey")
print(a)
b = foobar("hey")
print(b)```
I see, thanks !
This class is cool! Thanks!
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Another thing to be aware of with this implementation: the
wrappercontext manager is a sync function (not an async context manager, which do exist). So, what this means, I think, is that everything in your wrapper is executed in the main event loop, when this decorates an async function.If you made some database call, external http request, or file access, inside the decorator (
wrapperhere) then it could stall your event loop.I think that if you were doing this sort of thing in your decorator, you could possibly, inside of
call_async, run the context manager inside of a ThreadPoolExecutor to avoid locking the event loop.Of course, if you are okay with duplicating code, you could also use the async context manager instead, but in that case, I think I'd just skip the context manager.