dycw-utilities 0.162.7__tar.gz → 0.162.8__tar.gz

{dycw_utilities-0.162.7 → dycw_utilities-0.162.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.162.7
+Version: 0.162.8
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12

{dycw_utilities-0.162.7 → dycw_utilities-0.162.8}/pyproject.toml

@@ -230,7 +230,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.162.7"
+version = "0.162.8"
 
 [project.entry-points.pytest11]
 pytest-randomly = "utilities.pytest_plugins.pytest_randomly"
@@ -263,7 +263,7 @@ test = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.162.7"
+current_version = "0.162.8"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.162.7 → dycw_utilities-0.162.8}/src/tests/test_redis.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 
 from asyncio import Queue
 from itertools import chain, repeat
-from typing import TYPE_CHECKING, Any
+from logging import getLogger
+from typing import TYPE_CHECKING, Any, ClassVar, cast
 
 from hypothesis import HealthCheck, Phase, given, settings
 from hypothesis.strategies import (
@@ -15,18 +16,22 @@ from hypothesis.strategies import (
     permutations,
     sampled_from,
 )
-from pytest import mark, param, raises
+from pytest import LogCaptureFixture, fixture, mark, param, raises
 from redis.asyncio import Redis
 from redis.asyncio.client import PubSub
 
 from tests.conftest import SKIPIF_CI_AND_NOT_LINUX
 from tests.test_objects.objects import objects
 from utilities.asyncio import get_items_nowait, sleep_td
+from utilities.functions import get_class_name, identity
 from utilities.hypothesis import int64s, pairs, text_ascii
+from utilities.iterables import one
 from utilities.operator import is_equal
 from utilities.orjson import deserialize, serialize
 from utilities.redis import (
     PublishError,
+    _decoded_data,
+    _handle_message,
     _is_message,
     _RedisMessage,
     publish,
@@ -50,6 +55,98 @@ if TYPE_CHECKING:
 _PUB_SUB_SLEEP: TimeDelta = 0.1 * SECOND
 
 
+@fixture
+def queue() -> Queue[Any]:
+    return Queue()
+
+
+class TestHandleMessage:
+    message: ClassVar[_RedisMessage] = _RedisMessage(
+        type="message", pattern=None, channel=b"channel", data=b"data"
+    )
+
+    def test_main(self, *, queue: Queue[Any]) -> None:
+        _handle_message(self.message, identity, queue)
+        assert queue.qsize() == 1
+        assert queue.get_nowait() == self.message
+
+    def test_transform(self, *, queue: Queue[Any]) -> None:
+        _handle_message(self.message, _decoded_data, queue)
+        assert queue.qsize() == 1
+        assert queue.get_nowait() == "data"
+
+    def test_error_transform_no_handler(self, *, queue: Queue[Any]) -> None:
+        class CustomError(Exception): ...
+
+        def transform(message: _RedisMessage, /) -> None:
+            raise CustomError(message)
+
+        _handle_message(self.message, transform, queue)
+        assert queue.empty()
+
+    def test_error_transform_with_handler(
+        self, *, caplog: LogCaptureFixture, queue: Queue[Any]
+    ) -> None:
+        logger = getLogger(name := unique_str())
+
+        class CustomError(Exception): ...
+
+        def transform(message: _RedisMessage, /) -> None:
+            raise CustomError(message)
+
+        def error_transform(message: _RedisMessage, error: Exception, /) -> None:
+            logger.warning("Got %r transforming %r", get_class_name(error), message)
+
+        _handle_message(self.message, transform, queue, error_transform=error_transform)
+        assert queue.empty()
+        record = one(r for r in caplog.records if r.name == name)
+        assert record.message == f"Got 'CustomError' transforming {self.message}"
+
+    @mark.parametrize(
+        ("min_length", "expected"),
+        [param(3, False), param(4, False), param(5, True), param(6, True)],
+    )
+    def test_filter(
+        self, *, queue: Queue[Any], min_length: int, expected: bool
+    ) -> None:
+        _handle_message(
+            self.message,
+            _decoded_data,
+            queue,
+            filter_=lambda text: len(text) >= min_length,
+        )
+        result = queue.empty()
+        assert result is expected
+
+    def test_error_filter_no_handler(self, *, queue: Queue[Any]) -> None:
+        _handle_message(
+            self.message,
+            identity,
+            queue,
+            filter_=lambda text: cast("Any", text)["invalid"],
+        )
+        assert queue.empty()
+
+    def test_error_filter_with_handler(
+        self, *, caplog: LogCaptureFixture, queue: Queue[Any]
+    ) -> None:
+        logger = getLogger(name := unique_str())
+
+        def error_filter(message: _RedisMessage, error: Exception, /) -> None:
+            logger.warning("Got %r filtering %r", get_class_name(error), message)
+
+        _handle_message(
+            self.message,
+            identity,
+            queue,
+            filter_=lambda text: cast("Any", text)["invalid"],
+            error_filter=error_filter,
+        )
+        assert queue.empty()
+        record = one(r for r in caplog.records if r.name == name)
+        assert record.message == f"Got 'KeyError' filtering {self.message}"
+
+
 class TestIsMessage:
     @mark.parametrize(
         ("message", "channels", "expected"),

{dycw_utilities-0.162.7 → dycw_utilities-0.162.8}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.162.7"
+__version__ = "0.162.8"

{dycw_utilities-0.162.7 → dycw_utilities-0.162.8}/src/utilities/redis.py

@@ -639,9 +639,11 @@ def subscribe(
     /,
     *,
     timeout: Delta | None = _SUBSCRIBE_TIMEOUT,
-    sleep: Delta = _SUBSCRIBE_SLEEP,
     output: Literal["raw"],
-    filter_: Callable[[_RedisMessage], bool] | None = None,
+    error_transform: Callable[[_RedisMessage, Exception], None] | None = None,
+    filter_: Callable[[bytes], bool] | None = None,
+    error_filter: Callable[[bytes, Exception], None] | None = None,
+    sleep: Delta = _SUBSCRIBE_SLEEP,
 ) -> AsyncIterator[Task[None]]: ...
 @overload
 @enhanced_async_context_manager
@@ -652,9 +654,11 @@ def subscribe(
     /,
     *,
     timeout: Delta | None = _SUBSCRIBE_TIMEOUT,
-    sleep: Delta = _SUBSCRIBE_SLEEP,
     output: Literal["bytes"],
+    error_transform: Callable[[_RedisMessage, Exception], None] | None = None,
     filter_: Callable[[bytes], bool] | None = None,
+    error_filter: Callable[[bytes, Exception], None] | None = None,
+    sleep: Delta = _SUBSCRIBE_SLEEP,
 ) -> AsyncIterator[Task[None]]: ...
 @overload
 @enhanced_async_context_manager
@@ -665,9 +669,11 @@ def subscribe(
     /,
     *,
     timeout: Delta | None = _SUBSCRIBE_TIMEOUT,
-    sleep: Delta = _SUBSCRIBE_SLEEP,
     output: Literal["text"] = "text",
+    error_transform: Callable[[_RedisMessage, Exception], None] | None = None,
     filter_: Callable[[str], bool] | None = None,
+    error_filter: Callable[[str, Exception], None] | None = None,
+    sleep: Delta = _SUBSCRIBE_SLEEP,
 ) -> AsyncIterator[Task[None]]: ...
 @overload
 @enhanced_async_context_manager
@@ -678,34 +684,35 @@ def subscribe[T](
     /,
     *,
     timeout: Delta | None = _SUBSCRIBE_TIMEOUT,
-    sleep: Delta = _SUBSCRIBE_SLEEP,
     output: Callable[[bytes], T],
+    error_transform: Callable[[_RedisMessage, Exception], None] | None = None,
     filter_: Callable[[T], bool] | None = None,
+    error_filter: Callable[[T, Exception], None] | None = None,
+    sleep: Delta = _SUBSCRIBE_SLEEP,
 ) -> AsyncIterator[Task[None]]: ...
 @enhanced_async_context_manager
-async def subscribe[T](
+async def subscribe[T, U](
     redis: Redis,
     channels: MaybeIterable[str],
     queue: Queue[_RedisMessage] | Queue[bytes] | Queue[T],
     /,
     *,
     timeout: Delta | None = _SUBSCRIBE_TIMEOUT,
-    sleep: Delta = _SUBSCRIBE_SLEEP,
     output: Literal["raw", "bytes", "text"] | Callable[[bytes], T] = "text",
-    filter_: Callable[[Any], bool] | None = None,
+    error_transform: Callable[[_RedisMessage, Exception], None] | None = None,
+    filter_: Callable[[T], bool] | None = None,
+    error_filter: Callable[[T, Exception], None] | None = None,
+    sleep: Delta = _SUBSCRIBE_SLEEP,
 ) -> AsyncIterator[Task[None]]:
     """Subscribe to the data of a given channel(s)."""
     channels = list(always_iterable(channels))  # skipif-ci-and-not-linux
     match output:  # skipif-ci-and-not-linux
         case "raw":
-            transform = cast("Any", identity)
+            transform = cast("Callable[[_RedisMessage], T]", identity)
         case "bytes":
-            transform = cast("Any", itemgetter("data"))
+            transform = cast("Callable[[_RedisMessage], T]", itemgetter("data"))
         case "text":
-
-            def transform(message: _RedisMessage, /) -> str:  # pyright: ignore[reportRedeclaration]
-                return message["data"].decode()
-
+            transform = cast("Callable[[_RedisMessage], T]", _decoded_data)
         case Callable() as deserialize:
 
             def transform(message: _RedisMessage, /) -> T:
@@ -721,8 +728,10 @@ async def subscribe[T](
             transform,
             queue,
             timeout=timeout,
-            sleep=sleep,
+            error_transform=error_transform,
             filter_=filter_,
+            error_filter=error_filter,
+            sleep=sleep,
         )
     )
     try:  # skipif-ci-and-not-linux
@@ -737,16 +746,22 @@ async def subscribe[T](
             await task
 
 
-async def _subscribe_core(
+def _decoded_data(message: _RedisMessage, /) -> str:
+    return message["data"].decode()
+
+
+async def _subscribe_core[T](
     redis: Redis,
     channels: MaybeIterable[str],
-    transform: Callable[[_RedisMessage], Any],
+    transform: Callable[[_RedisMessage], T],
     queue: Queue[Any],
     /,
     *,
     timeout: Delta | None = _SUBSCRIBE_TIMEOUT,
+    error_transform: Callable[[_RedisMessage, Exception], None] | None = None,
+    filter_: Callable[[T], bool] | None = None,
+    error_filter: Callable[[T, Exception], None] | None = None,
     sleep: Delta = _SUBSCRIBE_SLEEP,
-    filter_: Callable[[Any], bool] | None = None,
 ) -> None:
     timeout_use = (  # skipif-ci-and-not-linux
         None if timeout is None else to_seconds(timeout)
@@ -758,11 +773,15 @@ async def _subscribe_core(
         while True:
             message = await pubsub.get_message(timeout=timeout_use)
             if is_subscribe_message(message):
-                transformed = transform(message)
-                if (filter_ is None) or filter_(transformed):
-                    queue.put_nowait(transformed)
-            else:
-                await sleep_td(sleep)
+                _handle_message(
+                    message,
+                    transform,
+                    queue,
+                    error_transform=error_transform,
+                    filter_=filter_,
+                    error_filter=error_filter,
+                )
+            await sleep_td(sleep)
 
 
 def _is_message(
@@ -781,6 +800,33 @@ def _is_message(
     )
 
 
+def _handle_message[T](
+    message: _RedisMessage,
+    transform: Callable[[_RedisMessage], T],
+    queue: Queue[Any],
+    /,
+    *,
+    error_transform: Callable[[_RedisMessage, Exception], None] | None = None,
+    filter_: Callable[[T], bool] | None = None,
+    error_filter: Callable[[T, Exception], None] | None = None,
+) -> None:
+    try:
+        transformed = transform(message)
+    except Exception as error:  # noqa: BLE001
+        if error_transform is not None:
+            error_transform(message, error)
+        return
+    if filter_ is None:
+        queue.put_nowait(transformed)
+        return
+    try:
+        if filter_(transformed):
+            queue.put_nowait(transformed)
+    except Exception as error:  # noqa: BLE001
+        if error_filter is not None:
+            error_filter(transformed, error)
+
+
 class _RedisMessage(TypedDict):
     type: Literal["subscribe", "psubscribe", "message", "pmessage"]
     pattern: str | None