Improve add_callback_after_call decorator to add return value as call…

…back func argument. Fixed typing

src/omnipy/data/model.py

@@ -731,13 +731,22 @@ def _convert_to_model_if_reasonable(self, args, name, ret):
 
     def __getattr__(self, attr: str) -> Any:
         contents_attr = self._getattr_from_contents_obj(attr)
+
         if _is_interactive_mode() and not self._is_non_omnipy_pydantic_model():
-            contents_holder = AttribHolder(self, 'contents', copy_attr=True)
+            contents_holder_context = AttribHolder(self, 'contents', copy_attr=True)
 
             contents_cls_attr = self._getattr_from_contents_cls(attr)
+
+            def _validate_contents(ret: Any):
+                self.validate_contents()
+                return ret
+
             if not isinstance(contents_cls_attr, property) and callable(contents_attr):
                 contents_attr = add_callback_after_call(
                     contents_attr, self.validate_contents, with_context=contents_holder)
+                contents_attr = add_callback_after_call(contents_attr,
+                                                        _validate_contents,
+                                                        contents_holder_context)
 
         return contents_attr
 

src/omnipy/util/decorators.py

@@ -1,34 +1,50 @@
 from contextlib import AbstractContextManager
-from typing import Any, Callable, ParamSpec, TypeVar
+from typing import Any, Callable, Concatenate, ParamSpec, TypeVar
 
 _DecoratedP = ParamSpec('_DecoratedP')
 _DecoratedR = TypeVar('_DecoratedR')
 _CallbackP = ParamSpec('_CallbackP')
+_CallbackR = TypeVar('_CallbackR')
+_ReturnT = TypeVar('_ReturnT')
 
+no_context = None
 
-def add_callback_after_call(func: Callable[_DecoratedP, _DecoratedR],
-                            callback_func: Callable[_CallbackP, None],
+
+def add_callback_after_call(decorated_func: Callable[_DecoratedP, _DecoratedR],
+                            callback_func: Callable[Concatenate[_ReturnT | None, _CallbackP],
+                                                    _CallbackR],
+                            with_context: AbstractContextManager | None,
                             *cb_args: _CallbackP.args,
-                            with_context: AbstractContextManager | None = None,
                             **cb_kwargs: _CallbackP.kwargs) -> Callable[_DecoratedP, _DecoratedR]:
-    class ValidateAfterCall(AbstractContextManager):
+    class CallbackAfterCall(AbstractContextManager):
+        def __init__(self, callback_func, *args: _DecoratedP.args, **kwargs: _DecoratedP.kwargs):
+            self._callback_func = callback_func
+            self._args = args
+            self._kwargs = kwargs
+            self.return_value: _ReturnT | None = None
+
         def __enter__(self):
-            ...
+            self.return_value = self._callback_func(*self._args, **self._kwargs)
+            return self
 
         def __exit__(self, exc_type, exc_val, exc_tb):
             ret = None
 
             if exc_val is None:
-                callback_func(*cb_args, **cb_kwargs)
+                self.return_value = callback_func(self.return_value, *cb_args, **cb_kwargs)
             return ret
 
+    def _callback_after_call(decorated_func, *args: _DecoratedP.args, **kwargs: _DecoratedP.kwargs):
+        with CallbackAfterCall(decorated_func, *args, **kwargs) as callback:
+            ...
+        return callback.return_value
+
     def _inner(*args: _DecoratedP.args, **kwargs: _DecoratedP.kwargs) -> _DecoratedR:
         if with_context:
             with with_context:
-                with ValidateAfterCall():
-                    return func(*args, **kwargs)
-        with ValidateAfterCall():
-            return func(*args, **kwargs)
+                return _callback_after_call(decorated_func, *args, **kwargs)
+
+        return _callback_after_call(decorated_func, *args, **kwargs)
 
     return _inner
 

tests/util/test_decorators.py

@@ -3,43 +3,47 @@
 import pytest
 
 from omnipy.util.contexts import AttribHolder
-from omnipy.util.decorators import add_callback_after_call, apply_decorator_to_property
+from omnipy.util.decorators import add_callback_after_call, apply_decorator_to_property, no_context
 
 
 def test_callback_after_func_call() -> None:
     def my_appender(a: list[int], b: int) -> list[int]:
         a.append(b)
         return a
 
-    def my_callback_after_call(x: list[int], *, y: int) -> None:
-        x.append(y)
+    def my_callback_after_call(ret: list[int] | None, x: int, *, y: int) -> list[int] | None:
+        if ret is not None:
+            ret.append(x)
+            ret.append(y)
+        return ret
 
     my_list = [1, 2, 3]
 
     decorated_my_appender = add_callback_after_call(
-        my_appender, my_callback_after_call, my_list, y=0)
+        my_appender, my_callback_after_call, no_context, 5, y=6)
 
     ret_list = decorated_my_appender(my_list, 4)
-    assert ret_list == my_list == [1, 2, 3, 4, 0]
+    assert ret_list == my_list == [1, 2, 3, 4, 5, 6]
 
 
 def test_callback_after_func_call_with_attrib_holder_error_in_func() -> None:
     class A:
         def __init__(self, numbers: list[int]) -> None:
             self.numbers = numbers
 
-    def my_appender(a: A, b: int) -> list[int]:
+    def my_appender(a: A, b: int) -> A:
         a.numbers.append(b)
         raise RuntimeError()
 
-    def my_callback_after_call(x: A, *, y: int) -> None:
+    def my_callback_after_call(ret: A | None, x: A, *, y: int) -> None:
+        assert ret is None
         x.numbers.append(y)
 
     my_a = A([1, 2, 3])
 
-    restore_numbers = AttribHolder(my_a, 'numbers', copy_attr=True)
+    restore_numbers_context = AttribHolder(my_a, 'numbers', copy_attr=True)
     decorated_my_appender = add_callback_after_call(
-        my_appender, my_callback_after_call, my_a, y=0, with_context=restore_numbers)
+        my_appender, my_callback_after_call, restore_numbers_context, my_a, y=0)
 
     try:
         decorated_my_appender(my_a, 4)
@@ -54,26 +58,30 @@ class A:
         def __init__(self, numbers: list[int]) -> None:
             self.numbers = numbers
 
-    def my_appender(a: A, b: int) -> list[int]:
+    def my_appender(a: A, b: int) -> A:
         a.numbers.append(b)
         return a
 
-    def my_callback_after_call(x: A, *, y: int) -> None:
+    def my_callback_after_call(ret: A | None, x: A, *, y: int) -> None:
+        if ret is not None:
+            ret.numbers.append(y)
         x.numbers.append(y)
         raise RuntimeError()
 
     my_a = A([1, 2, 3])
+    my_other_a = A([1, 2, 3])
 
-    restore_numbers = AttribHolder(my_a, 'numbers', copy_attr=True)
+    restore_numbers_context = AttribHolder(my_a, 'numbers', copy_attr=True)
     decorated_my_appender = add_callback_after_call(
-        my_appender, my_callback_after_call, my_a, y=0, with_context=restore_numbers)
+        my_appender, my_callback_after_call, restore_numbers_context, my_other_a, y=4)
 
     try:
         decorated_my_appender(my_a, 4)
     except RuntimeError:
         pass
 
     assert my_a.numbers == [1, 2, 3]
+    assert my_other_a.numbers == [1, 2, 3, 4]
 
 
 def test_apply_decorator_to_property():