queue-func 0.1__py3-none-any.whl

queue_func/__init__.py

@@ -0,0 +1 @@
+# coding: utf-8

queue_func/broker/__init__.py

@@ -0,0 +1 @@
+# coding: utf-8

queue_func/broker/base.py

@@ -0,0 +1,82 @@
+# coding: utf-8
+
+import abc
+import typing as t
+
+from queue_func.message.base import BaseMessage, Packed
+
+if t.TYPE_CHECKING:
+    from queue_func.config.general import Config
+    from queue_func.producer.base import BaseProducer
+    from queue_func.consumer.base import BaseConsumer
+
+
+ConfigType = t.TypeVar("ConfigType", bound="Config")
+
+
+class BaseBroker(t.Generic[ConfigType, Packed], metaclass=abc.ABCMeta):
+    """一条通道的装配者。
+
+    创建时准备好双方共用的 Message，并持有配置与 Producer/Consumer。
+    对外通过 publish、start_consumer 等入口操作；运行期业务逻辑仍在 Producer/Consumer 内完成。
+    """
+
+    __slots__ = (
+        "_config",
+        "_message",
+        "_producer",
+        "_consumer",
+    )
+
+    def __init__(self, config: ConfigType):
+
+        self._config = config
+        self._message: BaseMessage[Packed] = self._config.message_cls(
+            message_serializer_cls=self._config.message_serializer_cls,
+            enable_serializer=self._config.enable_message_serializer,
+        )
+        self._producer = self._create_producer()
+        self._consumer = self._create_consumer()
+
+    @abc.abstractmethod
+    def _create_producer(self) -> "BaseProducer[Packed]":
+
+        pass
+
+    @abc.abstractmethod
+    def _create_consumer(self) -> "BaseConsumer[Packed]":
+
+        pass
+
+    @property
+    def message(self) -> BaseMessage[Packed]:
+
+        return self._message
+
+    async def publish(
+            self,
+            coro_args: tuple[t.Any, ...],
+            coro_kwargs: dict[str, t.Any] | None = None,
+            **kwargs: t.Any) -> None:
+
+        await self._producer.publish(coro_args, coro_kwargs, **kwargs)
+
+    async def start_consumer(self) -> None:
+
+        await self._consumer.consume()
+
+    async def stop_consumer(self) -> None:
+
+        await self._consumer.stop()
+
+    async def join_consumer(self) -> None:
+
+        await self._consumer.join()
+
+    async def stop_producer(self) -> None:
+
+        await self._producer.stop()
+
+    async def join_producer(self) -> None:
+
+        await self._producer.join()

queue_func/broker/queue.py

@@ -0,0 +1,44 @@
+# coding: utf-8
+
+import asyncio
+import typing as t
+
+from queue_func.message.base import Packed
+from queue_func.broker.base import BaseBroker as _BaseBroker
+from queue_func.producer.queue import Producer
+from queue_func.consumer.queue import Consumer
+
+if t.TYPE_CHECKING:
+    from queue_func.config.queue import Config
+
+
+class Broker(_BaseBroker["Config", Packed]):
+    """内存队列通道。队列由这里创建，再交给 Producer 和 Consumer。"""
+
+    __slots__ = (
+        "_queue",
+    )
+
+    def __init__(self, config: "Config"):
+        self._queue: asyncio.Queue[Packed] = asyncio.Queue(maxsize=config.maxsize)
+        super().__init__(config)
+
+    def _create_producer(self) -> Producer[Packed]:
+
+        return Producer(
+            queue=self._queue,
+            publish_timeout=self._config.publish_timeout,
+            message=self._message,
+        )
+
+    def _create_consumer(self) -> Consumer[Packed]:
+
+        return Consumer(
+            queue=self._queue,
+            consume_timeout=self._config.consume_timeout,
+            message=self._message,
+            coro=self._config.coro,
+            enable_consumer=self._config.enable_consumer,
+            pool_executor_cls=self._config.consumer_pool_executor_cls,
+            pool_maxsize=self._config.consumer_pool_maxsize,
+        )

queue_func/config/__init__.py

@@ -0,0 +1 @@
+# coding: utf-8

queue_func/config/general.py

@@ -0,0 +1,32 @@
+# coding: utf-8
+
+import typing as t
+
+from pydantic.dataclasses import dataclass as pyd_dataclass
+
+from queue_func.config.utils import ExtraIgnoreConfig
+from queue_func.message.base import BaseMessage
+from queue_func.message.general import Message
+from queue_func.serializer.base import BaseSerializer
+from queue_func.serializer.orjson_ser import OrJSONSerializer
+from queue_func.pool_executor.base import BasePoolExecutor
+from queue_func.pool_executor.general import PoolExecutor
+
+
+@pyd_dataclass(frozen=True, config=ExtraIgnoreConfig)
+class Config:
+
+    # consumer
+    coro: t.Callable[[...], t.Coroutine]
+    enable_consumer: bool = True
+    consume_timeout: float = 2.0
+    consumer_pool_executor_cls: type[BasePoolExecutor] = PoolExecutor
+    consumer_pool_maxsize: int = 200
+
+    # broker
+    message_cls: type[BaseMessage] = Message
+    message_serializer_cls: type[BaseSerializer] = OrJSONSerializer
+    enable_message_serializer: bool = True
+
+    # producer
+    publish_timeout: float | None = None

queue_func/config/queue.py

@@ -0,0 +1,14 @@
+# coding: utf-8
+
+import typing as t
+from pydantic.dataclasses import dataclass as pyd_dataclass
+
+from queue_func.config.utils import ExtraIgnoreConfig
+from queue_func.config.general import Config as _Config
+
+
+@pyd_dataclass(frozen=True, config=ExtraIgnoreConfig)
+class Config(_Config):
+
+    # broker
+    maxsize: int = 200

queue_func/config/utils.py

@@ -0,0 +1,15 @@
+# coding: utf-8
+
+from pydantic.config import ConfigDict
+
+__all__ = (
+    "ExtraIgnoreConfig",
+)
+
+
+ExtraIgnoreConfig = ConfigDict(
+    extra="ignore",
+    validate_assignment=True,
+    arbitrary_types_allowed=True,
+    coerce_numbers_to_str=True,
+)

queue_func/consumer/__init__.py

@@ -0,0 +1 @@
+# coding: utf-8

queue_func/consumer/base.py

@@ -0,0 +1,97 @@
+# coding: utf-8
+
+import abc
+import asyncio
+import logging
+import contextlib
+import typing as t
+from functools import wraps
+
+from queue_func.message.base import BaseMessage, Packed
+from queue_func.pool_executor.base import BasePoolExecutor
+
+
+class BaseConsumer(t.Generic[Packed], metaclass=abc.ABCMeta):
+
+    __slots__ = (
+        "_message",
+        "_coro",
+        "_enable_consumer",
+        "_pool_executor",
+        "_task",
+        "_stop_flag",
+        "_logger",
+    )
+
+    def __init__(
+            self,
+            message: BaseMessage[Packed],
+            coro: t.Callable[..., t.Coroutine],
+            enable_consumer: bool,
+            pool_executor_cls: type[BasePoolExecutor],
+            pool_maxsize: int):
+
+        self._message = message
+        self._coro = coro
+        self._enable_consumer = enable_consumer
+        self._pool_executor = self._init_pool_executor(pool_executor_cls, pool_maxsize)
+        self._task: asyncio.Task | None = None
+        self._stop_flag = False
+        self._logger = logging.getLogger(__name__)
+
+    def _init_pool_executor(
+            self,
+            consumer_pool_executor_cls: type[BasePoolExecutor],
+            consumer_pool_maxsize: int) -> BasePoolExecutor | None:
+
+        if self._enable_consumer:
+            return consumer_pool_executor_cls(consumer_pool_maxsize)
+        return None
+
+    def _run_until_stopped(
+            self,
+            coro: t.Callable[..., t.Awaitable],
+            delay: float = 0.02) -> t.Callable[..., t.Coroutine]:
+
+        @wraps(coro)
+        async def wrapped(*args, **kwargs) -> None:
+
+            while not self._stop_flag:
+                try:
+                    await coro(*args, **kwargs)
+                except Exception as e:
+                    self._logger.exception(e)
+                await asyncio.sleep(delay)
+
+        return wrapped
+
+    async def _dispatch(self, coro_args: tuple[t.Any], coro_kwargs: dict[str, t.Any], **_kwargs) -> None:
+
+        assert self._pool_executor is not None, "consumer pool not initialized"
+        await self._pool_executor.submit(self._coro, coro_args, coro_kwargs)
+
+    async def consume(self) -> None:
+
+        if self._enable_consumer:
+            self._stop_flag = False
+            consume_loop = self._run_until_stopped(self._consume)
+            self._task = asyncio.create_task(consume_loop())
+        return None
+
+    @abc.abstractmethod
+    async def _consume(self) -> None:  # pragma: no cover
+
+        pass
+
+    async def stop(self) -> None:
+
+        self._stop_flag = True
+
+    async def join(self) -> None:
+
+        if self._task is not None:
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._task
+
+        if self._pool_executor is not None:
+            await self._pool_executor.join()

queue_func/consumer/queue.py

@@ -0,0 +1,55 @@
+# coding: utf-8
+
+import asyncio
+import typing as t
+
+from queue_func.message.base import BaseMessage, Packed
+from queue_func.pool_executor.base import BasePoolExecutor
+from queue_func.consumer.base import BaseConsumer as _BaseConsumer
+
+
+class Consumer(_BaseConsumer[Packed]):
+
+    __slots__ = (
+        "_queue",
+        "_consume_timeout",
+    )
+
+    def __init__(
+            self,
+            queue: asyncio.Queue[Packed],
+            consume_timeout: float,
+            message: BaseMessage[Packed],
+            coro: t.Callable[..., t.Coroutine],
+            enable_consumer: bool,
+            pool_executor_cls: type[BasePoolExecutor],
+            pool_maxsize: int):
+
+        super().__init__(
+            message,
+            coro,
+            enable_consumer,
+            pool_executor_cls,
+            pool_maxsize,
+        )
+        self._queue = queue
+        self._consume_timeout = consume_timeout
+
+    async def _consume(self) -> None:
+
+        while (not self._stop_flag) or (not self._queue.empty()):
+            try:
+                message: Packed = await asyncio.wait_for(
+                    self._queue.get(),
+                    timeout=self._consume_timeout,
+                )
+            except asyncio.TimeoutError:
+                continue
+            coro_args, coro_kwargs = self._message.unpack(message)
+            await self._dispatch(coro_args, coro_kwargs)
+            self._queue.task_done()
+
+    async def join(self) -> None:
+
+        await super().join()
+        await self._queue.join()

queue_func/factory/__init__.py

@@ -0,0 +1,12 @@
+# coding: utf-8
+
+from queue_func.factory.backend import Backend
+from queue_func.factory.deco import QueueFunc, queue
+from queue_func.factory.runtime import QueueFuncRuntime
+
+__all__ = (
+    "Backend",
+    "QueueFunc",
+    "QueueFuncRuntime",
+    "queue",
+)

queue_func/factory/backend.py

@@ -0,0 +1,13 @@
+# coding: utf-8
+
+from enum import StrEnum, unique
+
+__all__ = (
+    "Backend",
+)
+
+
+@unique
+class Backend(StrEnum):
+
+    QUEUE = "queue"

queue_func/factory/builders/__init__.py

@@ -0,0 +1,9 @@
+# coding: utf-8
+
+"""Import builder modules so they register with :func:`register_backend`."""
+
+from queue_func.factory.builders import queue as queue
+
+__all__ = (
+    "queue",
+)

queue_func/factory/builders/queue.py

@@ -0,0 +1,20 @@
+# coding: utf-8
+
+import typing as t
+
+from queue_func.factory.backend import Backend
+from queue_func.factory.registry import register_backend
+
+if t.TYPE_CHECKING:
+    from queue_func.broker.base import BaseBroker
+
+
+@register_backend(Backend.QUEUE)
+def build_queue_broker(
+        coro: t.Callable[..., t.Coroutine],
+        **kwargs: t.Any) -> "BaseBroker[t.Any, t.Any]":
+
+    from queue_func.broker.queue import Broker
+    from queue_func.config.queue import Config
+
+    return Broker(Config(coro=coro, **kwargs))

queue_func/factory/deco.py

@@ -0,0 +1,93 @@
+# coding: utf-8
+
+import typing as t
+
+from queue_func.factory import builders as _builders  # noqa: F401 — register backends
+from queue_func.factory.backend import Backend
+from queue_func.factory.registry import build_broker
+
+if t.TYPE_CHECKING:
+    from queue_func.broker.base import BaseBroker
+    from queue_func.config.general import Config as GeneralConfig
+    from queue_func.message.base import BaseMessage
+    from queue_func.pool_executor.base import BasePoolExecutor
+    from queue_func.serializer.base import BaseSerializer
+
+_P = t.ParamSpec("_P")
+_R = t.TypeVar("_R")
+
+__all__ = (
+    "QueueFunc",
+    "queue",
+)
+
+
+class QueueFunc(t.Generic[_P, _R]):
+    """可调用协程，并附带 push。"""
+
+    def __init__(
+            self,
+            fn: t.Callable[_P, t.Coroutine[t.Any, t.Any, _R]],
+            broker: "BaseBroker[t.Any, t.Any]") -> None:
+
+        self._fn = fn
+        self._broker = broker
+        self.__wrapped__ = fn
+        self.__name__ = fn.__name__
+        self.__qualname__ = fn.__qualname__
+        self.__doc__ = fn.__doc__
+        self.__annotations__ = getattr(fn, "__annotations__", {})
+
+    async def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R:
+
+        return await self._fn(*args, **kwargs)
+
+    async def push(
+            self,
+            *,
+            args: tuple[t.Any, ...] = (),
+            kwargs: dict[str, t.Any] | None = None) -> None:
+
+        await self._broker.publish(coro_args=args, coro_kwargs=kwargs)
+
+    @property
+    def broker(self) -> "BaseBroker[t.Any, t.Any]":
+
+        return self._broker
+
+
+def _queue_impl(
+        backend: Backend,
+        **config_kwargs: t.Any) -> t.Callable[[t.Callable[_P, t.Coroutine[t.Any, t.Any, _R]]], QueueFunc[_P, _R]]:
+
+    def decorator(fn: t.Callable[_P, t.Coroutine[t.Any, t.Any, _R]]) -> QueueFunc[_P, _R]:
+
+        broker = build_broker(backend, fn, **config_kwargs)
+        return QueueFunc(fn, broker)
+
+    return decorator
+
+
+@t.overload
+def queue(
+        backend: t.Literal[Backend.QUEUE],
+        *,
+        enable_consumer: bool = ...,
+        consume_timeout: float = ...,
+        consumer_pool_executor_cls: type["BasePoolExecutor"] = ...,
+        consumer_pool_maxsize: int = ...,
+        message_cls: type["BaseMessage"] = ...,
+        message_serializer_cls: type["BaseSerializer"] = ...,
+        enable_message_serializer: bool = ...,
+        publish_timeout: float | None = ...,
+        maxsize: int = ...,
+) -> t.Callable[[t.Callable[_P, t.Coroutine[t.Any, t.Any, _R]]], QueueFunc[_P, _R]]:
+
+    ...
+
+
+def queue(
+        backend: Backend,
+        **config_kwargs: t.Any) -> t.Callable[[t.Callable[_P, t.Coroutine[t.Any, t.Any, _R]]], QueueFunc[_P, _R]]:
+
+    return _queue_impl(backend, **config_kwargs)

queue_func/factory/registry.py

@@ -0,0 +1,41 @@
+# coding: utf-8
+
+import typing as t
+
+from queue_func.factory.backend import Backend
+
+if t.TYPE_CHECKING:
+    from queue_func.broker.base import BaseBroker
+
+__all__ = (
+    "register_backend",
+    "build_broker",
+)
+
+BrokerBuilder = t.Callable[..., "BaseBroker[t.Any, t.Any]"]
+
+_BUILDERS: dict[Backend, BrokerBuilder] = {}
+
+
+def register_backend(backend: Backend) -> t.Callable[[BrokerBuilder], BrokerBuilder]:
+
+    def decorator(fn: BrokerBuilder) -> BrokerBuilder:
+
+        if backend in _BUILDERS:
+            raise ValueError(f"backend {backend!r} is already registered")
+        _BUILDERS[backend] = fn
+        return fn
+
+    return decorator
+
+
+def build_broker(
+        backend: Backend,
+        coro: t.Callable[..., t.Coroutine],
+        **kwargs: t.Any) -> "BaseBroker[t.Any, t.Any]":
+
+    try:
+        builder = _BUILDERS[backend]
+    except KeyError:
+        raise NotImplementedError(f"backend {backend!r} is not implemented yet") from None
+    return builder(coro, **kwargs)

queue_func/factory/runtime.py

@@ -0,0 +1,40 @@
+# coding: utf-8
+
+import typing as t
+
+from queue_func.factory.deco import QueueFunc
+
+__all__ = (
+    "QueueFuncRuntime",
+)
+
+
+class QueueFuncRuntime:
+    """按固定顺序批量启动 / 关闭多个 :class:`QueueFunc`。
+
+    stop：consumer 逐个 stop → join；producer 先全体 stop，再全体 join。
+    """
+
+    __slots__ = (
+        "_funcs",
+    )
+
+    def __init__(self, funcs: t.Iterable[QueueFunc[t.Any, t.Any]]) -> None:
+
+        self._funcs = list(funcs)
+
+    async def start(self) -> None:
+
+        for func in reversed(self._funcs):
+            await func.broker.start_consumer()
+
+    async def stop(self) -> None:
+
+        for func in self._funcs:
+            await func.broker.stop_consumer()
+            await func.broker.join_consumer()
+
+        for func in self._funcs:
+            await func.broker.stop_producer()
+        for func in self._funcs:
+            await func.broker.join_producer()

queue_func/message/__init__.py

@@ -0,0 +1 @@
+# coding: utf-8

queue_func/message/base.py

@@ -0,0 +1,39 @@
+# coding: utf-8
+
+import abc
+import typing as t
+
+from queue_func.serializer.base import BaseSerializer
+
+Packed = t.TypeVar("Packed")
+
+type ArgsMessage = tuple[tuple[t.Any], dict[str, t.Any]]
+
+
+class BaseMessage(t.Generic[Packed], metaclass=abc.ABCMeta):
+
+    __slots__ = (
+        "_enable_serializer",
+        "_serializer",
+    )
+
+    def __init__(
+            self,
+            message_serializer_cls: type[BaseSerializer],
+            enable_serializer: bool):
+
+        self._serializer = message_serializer_cls()
+        self._enable_serializer = enable_serializer
+
+    @abc.abstractmethod
+    def pack(
+            self,
+            coro_args: tuple[t.Any],
+            coro_kwargs: dict[str, t.Any] | None = None) -> Packed:
+
+        pass
+
+    @abc.abstractmethod
+    def unpack(self, message: Packed) -> ArgsMessage:
+
+        pass

queue_func/message/general.py

@@ -0,0 +1,30 @@
+# coding: utf-8
+
+import typing as t
+
+from queue_func.message.base import ArgsMessage, BaseMessage
+
+type PackedMessage = ArgsMessage | bytes
+
+
+class Message(BaseMessage[PackedMessage]):
+
+    __slots__ = ()
+
+    def pack(
+            self,
+            coro_args: tuple[t.Any],
+            coro_kwargs: dict[str, t.Any] | None = None) -> PackedMessage:
+
+        if coro_kwargs is None:
+            coro_kwargs = {}
+        message: ArgsMessage = (coro_args, coro_kwargs)
+        if self._enable_serializer:
+            return self._serializer.dumps(message, return_bytes=True)
+        return message
+
+    def unpack(self, message: PackedMessage) -> ArgsMessage:
+
+        if self._enable_serializer:
+            return t.cast(ArgsMessage, self._serializer.loads(message))
+        return t.cast(ArgsMessage, message)

queue_func/pool_executor/__init__.py

@@ -0,0 +1 @@
+# coding: utf-8