File size: 987 Bytes
5374a2d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
import threading
from functools import wraps
from contextlib import nullcontext
def atomic(lock=None):
"""
threading safe decorator, it can be used to decorate a function or receive a lock:
1. directly decorate a function: @atomic
2. receive a lock: @atomic(lock=shared_lock)
"""
lock = lock or threading.Lock()
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
with lock:
return func(*args, **kwargs)
return wrapper
return decorator if not callable(lock) else decorator(lock)
def atomic_method(func):
"""
threading safe decorator for class methods.
If there are self._lock in the instance, it will use the lock. Otherwise, use nullcontext for execution.
"""
@wraps(func)
def wrapper(self, *args, **kwargs):
context = getattr(self, "_lock", nullcontext())
with context:
return func(self, *args, **kwargs)
return wrapper
|