Merge pull request #60 from mutating/develop

0.0.43

Author: pomponchik

cantok/tokens/abstract/abstract_token.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from threading import RLock
 from time import sleep
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 from printo import describe_call, not_none
 
@@ -71,7 +71,7 @@ def __repr__(self) -> str:
         cancelled = report.from_token is self and report.cause == CancelCause.CANCELLED
 
         return describe_call(
-            type(self).__name__,
+            type(self),
             [superpower, *self._tokens],
             {
                 'cancelled': cancelled,
@@ -222,7 +222,7 @@ def cancel(self) -> 'AbstractToken':
         self._cancelled = True
         return self
 
-    def check(self) -> None:
+    def check(self, *, exception: Optional[Union[Type[BaseException], BaseException]] = None) -> None:
         """
         Raises an exception if the token is cancelled; does nothing otherwise.
 
@@ -231,19 +231,47 @@ def check(self) -> None:
         - Automatic cancellation by a specific token type raises the corresponding
           subclass (e.g. TimeoutCancellationError for TimeoutToken).
 
+        :param exception: Optional keyword-only exception class or instance used
+                          instead of the standard cancellation exception. A class
+                          receives the standard cancellation message; an instance
+                          is raised as is.
+        :raises ValueError: If exception is neither an exception class nor instance.
+
         >>> token = SimpleToken()
         >>> token.check()   # nothing happens
         >>> token.cancel()
         >>> token.check()   # raises CancellationError
         """
+        if exception is not None and not (
+            isinstance(exception, BaseException)
+            or (isinstance(exception, type) and issubclass(exception, BaseException))
+        ):
+            raise ValueError('Only an exception instance or an exception class can be passed.')
+
         with self._lock:
             report = self._get_report()
 
-            if report.cause == CancelCause.CANCELLED:
-                report.from_token._raise_cancelled_exception()
+            if exception is None:
+                if report.cause == CancelCause.CANCELLED:
+                    report.from_token._raise_cancelled_exception()
+
+                elif report.cause == CancelCause.SUPERPOWER:
+                    report.from_token._raise_superpower_exception()
+
+            elif report.cause != CancelCause.NOT_CANCELLED:
+                if isinstance(exception, BaseException):
+                    raise exception
+
+                message = (
+                    'The token has been cancelled.'
+                    if report.cause == CancelCause.CANCELLED
+                    else report.from_token._get_superpower_exception_message()
+                ) + report.from_token._get_exception_message_suffix()
+
+                if issubclass(exception, CancellationError):
+                    raise exception(message, report.from_token)
 
-            elif report.cause == CancelCause.SUPERPOWER:
-                report.from_token._raise_superpower_exception()
+                raise exception(message)
 
     def _get_report(self, direct: bool = True) -> CancellationReport:
         if self._cancelled:

docs/what_are_tokens/exceptions.md

@@ -32,6 +32,18 @@ Each type of token (except [`DefaultToken`](../types_of_tokens/DefaultToken.md))
 
 When you call the `check()` method on any token, one of two things will happen. If it (or any of the tokens nested in it) has been cancelled by calling the `cancel()` method, `CancellationError` will always be raised. But if the cancellation occurred as a result of the unique ability of the token, such as timeout expiration for `TimeoutToken`, then an exception specific to this type of token will be raised.
 
+`check()` also accepts a keyword-only `exception` argument. Omit it or pass `None` to keep the behavior described above. Pass an exception class, and `check()` raises an instance of it with the standard message for the cancellation cause; pass an existing exception object, and `check()` raises that object as is. If the token is not cancelled, `check()` still does nothing and the override is not used.
+
+```python
+from cantok import SimpleToken
+
+token = SimpleToken()
+token.cancel()
+token.check(exception=RuntimeError)
+#> ...
+#> RuntimeError: The token has been cancelled.
+```
+
 `ConditionCancellationError`, `TimeoutCancellationError`, and `CounterCancellationError` are inherited from `CancellationError`, so if you're not sure which specific exception you're catching, catch `CancellationError`. All of the listed exceptions can also be imported separately:
 
 ```python

pyproject.toml

@@ -4,14 +4,13 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cantok"
-version = "0.0.42"
+version = "0.0.43"
 authors = [{ name = "Evgeniy Blinov", email = "zheni-b@yandex.ru" }]
 description = 'Implementation of the "Cancellation Token" pattern'
 readme = "README.md"
 requires-python = ">=3.8"
 dependencies = [
-    'printo>=0.0.28',
-    'transfunctions>=0.0.12',
+    'printo>=0.0.29',
     'typing_extensions>=4.5.0 ; python_version < "3.13"',
 ]
 classifiers = [

tests/typing/tokens/abstract/test_abstract_token.py

@@ -7,18 +7,53 @@
     from typing import assert_type
 
 import pytest
+from full_match import match
 
 from cantok import AbstractToken, SimpleToken
 
 
 @pytest.mark.mypy_testing
-def test_abstract_token_exposes_doc_attribute_as_optional_string():
+def test_abstract_token_exposes_doc_and_check_type_contracts():
     """
-    `AbstractToken` exposes `doc` as an optional string in the type contract.
+    Expose the public type contracts of `doc` and `check()`.
 
-    A concrete token with `doc` is assigned to an `AbstractToken` variable, and
-    `assert_type` fixes the public attribute type.
+    `doc` is `Optional[str]`. The keyword-only `exception` may be omitted or set
+    to `None`, and accepts classes and instances throughout the `BaseException`
+    hierarchy. Every `check()` call has inferred type `None`.
     """
     token: AbstractToken = SimpleToken(doc='d')
 
     assert_type(token.doc, Optional[str])
+    assert_type(token.check(), None)
+    assert_type(token.check(exception=None), None)
+    assert_type(token.check(exception=RuntimeError), None)
+    assert_type(token.check(exception=RuntimeError('custom message')), None)
+    assert_type(token.check(exception=SystemExit), None)
+    assert_type(token.check(exception=SystemExit('custom message')), None)
+
+
+@pytest.mark.mypy_testing
+def test_abstract_token_rejects_invalid_check_exception_types():
+    """
+    Reject non-exception values and classes statically and at runtime.
+
+    At runtime, both invalid forms raise the documented `ValueError` with its
+    exact message.
+    """
+    token: AbstractToken = SimpleToken()
+    expected_message = match('Only an exception instance or an exception class can be passed.')
+
+    with pytest.raises(ValueError, match=expected_message):
+        token.check(exception=1)  # E: [arg-type]
+
+    with pytest.raises(ValueError, match=expected_message):
+        token.check(exception=object)  # E: [arg-type]
+
+
+@pytest.mark.mypy_testing
+def test_abstract_token_check_exception_is_keyword_only():
+    """Reject positional exception overrides statically and at runtime."""
+    token: AbstractToken = SimpleToken()
+
+    with pytest.raises(TypeError):
+        token.check(ValueError)  # E: [misc]