Repository: darksidecat/aiopriman

Issue:
Add decorator for locking functions

Discussion:

Patch:
diff --git a/aiopriman/manager/utils.py b/aiopriman/manager/utils.py
index 78bbaa6..392f195 100644
--- a/aiopriman/manager/utils.py
+++ b/aiopriman/manager/utils.py
@@ -1,8 +1,10 @@
 from abc import abstractmethod
+from functools import wraps
 from types import TracebackType
-from typing import Optional, Type
+from typing import Any, AsyncContextManager, Callable, Optional, Type
 
 from aiopriman.sync_primitives import SyncPrimitive
+from aiopriman.utils.common import inspect_params
 
 
 class _ContextManagerMixin:
@@ -22,3 +24,32 @@ class _ContextManagerMixin:
     @abstractmethod
     def release(self, from_context_manager: bool) -> None:
         """Release synchronization primitive"""
+
+
+def lock(
+        manager: Callable[..., AsyncContextManager[Any]],
+        **dec_kwargs: Any
+) -> Callable[..., Any]:
+    def decorator(
+            func: Callable[..., Any]
+    ) -> Callable[..., Any]:
+
+        @wraps(func)
+        async def wrapped(
+                *args: Any,
+                **kwargs: Any
+        ) -> Any:
+            dec_storage_data = dec_kwargs.get('storage_data')
+            func_storage_data = kwargs.get('storage_data')
+            if dec_storage_data is None and func_storage_data is None:
+                raise ValueError("decorated function need keyword param storage_data")
+
+            kwargs = {**dec_kwargs, **kwargs}
+            man_params = inspect_params(manager, **kwargs)
+            func_params = inspect_params(func, **kwargs)
+
+            async with manager(**man_params):
+                return await func(*args, **func_params)
+
+        return wrapped
+    return decorator
diff --git a/aiopriman/utils/common.py b/aiopriman/utils/common.py
new file mode 100644
index 0000000..5192202
--- /dev/null
+++ b/aiopriman/utils/common.py
@@ -0,0 +1,13 @@
+import inspect
+from typing import Any, Dict
+
+
+def inspect_params(obj: Any, **kwargs: Any) -> Dict[str, Any]:
+    payload = {}
+    params = inspect.signature(obj).parameters.keys()
+
+    for k, v in kwargs.items():
+        if k in params:
+            payload[k] = v
+
+    return payload
diff --git a/examples/basic_usage.py b/examples/basic_usage.py
deleted file mode 100644
index e69de29..0000000
diff --git a/examples/decorator.py b/examples/decorator.py
new file mode 100644
index 0000000..2366f40
--- /dev/null
+++ b/examples/decorator.py
@@ -0,0 +1,43 @@
+import asyncio
+import logging
+
+from aiopriman.manager import LockManager, SemaphoreManager
+from aiopriman.manager.utils import lock
+from aiopriman.storage import StorageData
+
+storage_data = StorageData()
+
+
+@lock(LockManager, key="test2")  # will be used key value from function call
+async def run_lock(name):
+    logging.debug(f"HERE LOCKED {name}")
+    await asyncio.sleep(3)
+
+
+@lock(SemaphoreManager, storage_data=storage_data, key="test", value=5)  # using decorator params
+async def run_sem(name):
+    logging.debug(f"HERE SEM LOCKED {name}")
+    await asyncio.sleep(3)
+
+
+async def main_run(*args):
+    await asyncio.gather(*args)
+
+
+if __name__ == '__main__':
+    logging.basicConfig(
+        level=logging.DEBUG,
+        format='%(levelname)s:%(name)s:(%(filename)s).%(funcName)s(%(lineno)d):%(message)s'
+    )
+
+    tasks = []
+    storage_data = StorageData()
+    for i in range(1, 10):
+        tasks.append(run_lock(i, storage_data=storage_data,  key="test"))
+        tasks.append(run_sem(i))
+
+    asyncio.run(
+        main_run(
+           *tasks
+        )
+    )


Test patch:
diff --git a/tests/manager/test_utils.py b/tests/manager/test_utils.py
new file mode 100644
index 0000000..635b8a9
--- /dev/null
+++ b/tests/manager/test_utils.py
@@ -0,0 +1,49 @@
+import pytest
+
+from aiopriman.manager import LockManager
+from aiopriman.manager.utils import lock
+from aiopriman.storage import StorageData
+
+
+@pytest.mark.asyncio
+async def test_lock_raise_need_keyword():
+    @lock(LockManager)
+    async def func():
+        pass
+
+    with pytest.raises(ValueError):
+        await func()
+
+
+@pytest.mark.asyncio
+async def test_lock(storage_data):
+    @lock(LockManager)
+    async def func():
+        assert 'LockStorage:test' in storage_data
+        assert storage_data['LockStorage:test'].lock.locked()
+
+    await func(storage_data=storage_data, key="test")
+
+
+@pytest.mark.asyncio
+async def test_lock_func_params(storage_data):
+    storage_data_func = storage_data
+    storage_data_dec = StorageData()
+
+    @lock(LockManager, storage_data=storage_data_dec, key="test2")
+    async def func(storage_data):
+        assert 'LockStorage:test' in storage_data
+        assert storage_data is storage_data_func
+
+    await func(storage_data=storage_data, key="test")
+
+
+@pytest.mark.asyncio
+async def test_lock_dec_params(storage_data):
+    storage_data_dec = StorageData()
+
+    @lock(LockManager, storage_data=storage_data_dec, key="test2")
+    async def func():
+        assert 'LockStorage:test2' in storage_data_dec
+
+    await func()
diff --git a/tests/utils/test_common.py b/tests/utils/test_common.py
new file mode 100644
index 0000000..c752034
--- /dev/null
+++ b/tests/utils/test_common.py
@@ -0,0 +1,8 @@
+from aiopriman.utils.common import inspect_params
+
+
+def test_inspect_params():
+    def func(a, b):
+        pass
+    kwargs = {'a': 1, 'b': 2, 'c': 3}
+    assert inspect_params(func, **kwargs) == {'a': 1, 'b': 2}