Create helper class decorator

Author: febus982

src/common/utils.py

@@ -0,0 +1,27 @@
+def apply_decorator_to_methods(
+    decorator, protected_methods: bool = False, private_methods: bool = False
+):
+    """
+    Class decorator to apply a given function or coroutine decorator
+    to all functions and coroutines within a class.
+    """
+
+    def class_decorator(cls):
+        for attr_name, attr_value in cls.__dict__.items():
+            # Check if the attribute is a callable (method or coroutine)
+            if not callable(attr_value):
+                continue
+
+            if attr_name.startswith(f"_{cls.__name__}__"):
+                if not private_methods:
+                    continue
+
+            elif attr_name.startswith("_"):
+                if not protected_methods:
+                    continue
+
+            # Replace the original callable with the decorated version
+            setattr(cls, attr_name, decorator(attr_value))
+        return cls
+
+    return class_decorator

tests/common/__init__.py

@@ -0,0 +1,67 @@
+import asyncio
+
+import pytest
+
+from common.utils import apply_decorator_to_methods
+
+
+@pytest.mark.parametrize(
+    "apply_to_protected_methods",
+    [
+        True,
+        False,
+    ],
+)
+@pytest.mark.parametrize(
+    "apply_to_private_methods",
+    [
+        True,
+        False,
+    ],
+)
+async def test_class_decorator(
+    apply_to_protected_methods: bool,
+    apply_to_private_methods: bool,
+):
+    def add_ten_decorator(func):
+        def wrapper(*args, **kwargs):
+            result = func(*args, **kwargs)
+            return result + 10
+
+        async def async_wrapper(*args, **kwargs):
+            result = await func(*args, **kwargs)
+            return result + 10
+
+        return wrapper if not asyncio.iscoroutinefunction(func) else async_wrapper
+
+    @apply_decorator_to_methods(
+        decorator=add_ten_decorator,
+        protected_methods=apply_to_protected_methods,
+        private_methods=apply_to_private_methods,
+    )
+    class MyClass:
+        def get_public(self):
+            return 10
+
+        def _get_protected(self):
+            return 10
+
+        def __get_private(self):
+            return 10
+
+        async def get_apublic(self):
+            return 10
+
+        async def _get_aprotected(self):
+            return 10
+
+        async def __get_aprivate(self):
+            return 10
+
+    c = MyClass()
+    assert c.get_public() == 20
+    assert c._get_protected() == 20 if apply_to_protected_methods else 10
+    assert c._MyClass__get_private() == 20 if apply_to_private_methods else 10
+    assert await c.get_apublic() == 20
+    assert await c._get_aprotected() == 20 if apply_to_protected_methods else 10
+    assert await c._MyClass__get_aprivate() == 20 if apply_to_private_methods else 10