improve assert_or_throw (#68)

* improve assert_or_throw

* update

README.md

@@ -19,6 +19,10 @@ pip install triad
 
 ## Release History
 
+### 0.5.3
+
+* Lazy evaluation for `assert_or_throw`
+
 ### 0.5.2
 
 * For pyarrow data conversion, support np.ndarray -> list

tests/utils/test_assertion.py

@@ -1,6 +1,8 @@
-from triad.utils.assertion import assert_or_throw, assert_arg_not_none
+import pickle
+
 from pytest import raises
 from triad.exceptions import NoneArgumentError
+from triad.utils.assertion import assert_arg_not_none, assert_or_throw
 
 
 def test_assert_or_throw():
@@ -9,6 +11,23 @@ def test_assert_or_throw():
     raises(AssertionError, lambda: assert_or_throw(False, "test"))
     raises(AssertionError, lambda: assert_or_throw(False, 123))
     raises(TypeError, lambda: assert_or_throw(False, TypeError()))
+
+    # lazy evaluation
+    raises(TypeError, lambda: assert_or_throw(False, lambda: TypeError()))
+
+    def fail_without_lazy():
+        raise TypeError
+
+    assert_or_throw(True, fail_without_lazy)
+    raises(TypeError, lambda: assert_or_throw(False, fail_without_lazy))
+
+    # serialization + lazy
+    mock = pickle.loads(pickle.dumps(Mock()))
+    raises(TypeError, lambda: mock.t1(False))
+    mock.t1(True)
+
+    raises(NotImplementedError, lambda: mock.t2(False))
+    mock.t2(True)
 
 
 def test_assert_arg_not_none():
@@ -25,3 +44,14 @@ def test_assert_arg_not_none():
     with raises(NoneArgumentError) as err:
         assert_arg_not_none(None, None, msg="b")
     assert "b" == err.value.args[0]
+
+
+class Mock:
+    def t1(self, v):
+        def fail_without_lazy():
+            raise TypeError
+
+        assert_or_throw(v, fail_without_lazy)
+
+    def t2(self, v):
+        assert_or_throw(v, lambda: NotImplementedError())

triad/utils/assertion.py

@@ -1,24 +1,48 @@
-from typing import Union, Any
+from typing import Any
 from triad.exceptions import NoneArgumentError
 
 
-def assert_or_throw(
-    bool_exp: bool, exception: Union[None, str, Exception] = None
-) -> None:
+def assert_or_throw(bool_exp: bool, exception: Any = None) -> None:
     """Assert on expression and throw custom exception
 
     :param bool_exp: boolean expression to assert on
     :param exception: a custom Exception instance, or any other object that
-        will be stringfied and instantiate an AssertionError
+        will be stringfied and instantiate an AssertionError, or a function
+        that can generate the supported data types
+
+    :Examples:
+
+        .. code-block:: python
+
+            assert_or_throw(True, "assertion error")
+            assert_or_throw(False)  # raise AssertionError
+            assert_or_throw(False, "assertion error")  # raise AssertionError
+            assert_or_throw(False, TypeError("assertion error"))  # raise TypeError
+
+            # Lazy evaluations is useful when constructing the error
+            # itself is expensive or error-prone. With lazy evaluations, happy
+            # path will be fast and error free.
+            def fail():  # a function that is slow and wrong
+                sleep(10)
+                raise TypeError
+
+            assert_or_throw(True, fail())  # (unexpectedly) raise TypeError
+            assert_or_throw(True, fail)  # no exception
+            assert_or_throw(True, lambda: "a" + fail())  # no exception
+            assert_or_throw(False, lambda: "a" + fail())  # raise TypeError
+
     """
     if not bool_exp:
-        if isinstance(exception, Exception):
-            raise exception
-        if isinstance(exception, str):
-            raise AssertionError(exception)
-        if exception is None:
+        _exception: Any = exception
+        if callable(exception):
+            _exception = exception()
+        if _exception is None:
             raise AssertionError()
-        raise AssertionError(str(exception))
+        if isinstance(_exception, Exception):
+            raise _exception
+        if isinstance(_exception, str):
+            raise AssertionError(_exception)
+        raise AssertionError(str(_exception))
 
 
 def assert_arg_not_none(obj: Any, arg_name: str = "", msg: str = "") -> None:

triad/utils/convert.py

@@ -36,6 +36,7 @@ def get_caller_global_local_vars(
     :Examples:
 
         .. code-block:: python
+
             def caller():
                 x=1
                 assert 1 == get_value("x")
@@ -57,6 +58,7 @@ def get_value(var_name):
     :Examples:
 
         .. code-block:: python
+
             def f1():
                 x=1
 
@@ -119,6 +121,7 @@ def str_to_object(
     :Examples:
 
         .. code-block:: python
+
             class _Mock(object):
                 def __init__(self, x=1):
                     self.x = x

triad_version/__init__.py

@@ -1 +1 @@
-__version__ = "0.5.2"
+__version__ = "0.5.3"