kuu 0.0.1__tar.gz

kuu-0.0.1/PKG-INFO

@@ -0,0 +1,57 @@
+Metadata-Version: 2.4
+Name: kuu
+Version: 0.0.1
+Summary: lightweight async distributed task queue
+Keywords: apscheduler,message queue,queue,schedule tasks,scheduler,scheduling,task,task queue
+Author: Alexey Pechenin
+Author-email: Alexey Pechenin <me@pyrorhythm.dev>
+License-Expression: Apache-2.0
+Classifier: Development Status :: 4 - Beta
+Classifier: Framework :: AsyncIO
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Internet
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: System :: Distributed Computing
+Classifier: Topic :: System :: Networking
+Classifier: Typing :: Typed
+Requires-Dist: anyio>=4.10,<5
+Requires-Dist: croniter>=6.2
+Requires-Dist: orjson>=3.10
+Requires-Dist: pydantic>=2.12,<3
+Requires-Dist: typer>=0.24.2
+Requires-Dist: uvloop>=0.22.1
+Requires-Dist: watchfiles>=1.1.1
+Requires-Dist: msgspec>=0.21 ; extra == 'msgspec'
+Requires-Dist: nats-py>=2.7 ; extra == 'nats'
+Requires-Dist: prometheus-client>=0.20 ; extra == 'prometheus'
+Requires-Dist: redis>=6.0,<8 ; extra == 'redis'
+Requires-Python: >=3.12
+Project-URL: homepage, https://github.com/pyrorhythm/kuu
+Provides-Extra: msgspec
+Provides-Extra: nats
+Provides-Extra: prometheus
+Provides-Extra: redis
+Description-Content-Type: text/markdown
+
+# kuu
+
+```shell
+uv install kuu --prerelease=allow
+
+# extras availible: msgspec, nats, prometheus, redis
+```
+
+a native distributed queue that is rather simple and easy-to-integrate in production
+has task queue, redis xstream, nats jetstream support, result backend, middlewares and event/signals
+also has a built-in scheduler yk
+
+nothing serious in this project, just got tired of taskiq's undefined behaviour and `logging.getLogger("root")`
+maybe will build a webui for it idk

kuu-0.0.1/kuu/__init__.py

@@ -0,0 +1,29 @@
+from .app import Kuu
+from .brokers.base import Broker, Delivery
+from .context import Context
+from .exceptions import NotConnected, RejectErr, RetryErr, TaskError
+from .message import Message
+from .middleware.base import Middleware
+from .results.base import Result, ResultBackend
+from .serializers import JSONSerializer, Serializer
+from .handle import TaskHandle
+from .task import Task
+
+__all__ = [
+	"Kuu",
+	"Task",
+	"TaskHandle",
+	"Message",
+	"Context",
+	"RetryErr",
+	"RejectErr",
+	"TaskError",
+	"NotConnected",
+	"Middleware",
+	"Broker",
+	"Delivery",
+	"Result",
+	"ResultBackend",
+	"Serializer",
+	"JSONSerializer",
+]

kuu-0.0.1/kuu/_import.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import logging
+import os
+import sys
+from collections.abc import Generator, Sequence
+from contextlib import contextmanager
+from importlib import import_module
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def add_cwd_in_path() -> Generator[None]:
+	"""
+	adds current directory in python path
+
+	this context manager adds current directory in sys.path,
+	so all python files are discoverable now, without installing
+	current project
+
+	:yield none
+	"""
+	cwd = Path.cwd()
+	if str(cwd) in sys.path:
+		yield
+	else:
+		logger.debug(f"inserting {cwd} in sys.path")
+		sys.path.insert(0, str(cwd))
+		try:
+			yield
+		finally:
+			try:
+				sys.path.remove(str(cwd))
+			except ValueError:
+				logger.warning(f"cannot remove '{cwd}' from sys.path")
+
+
+def import_object(object_spec: str, app_dir: str | None = None) -> Any:
+	"""
+	it parses python object spec and imports it
+
+	:param object_spec: string in format like `package.module:variable`
+	:param app_dir: directory to add in sys.path for importing
+	:raises ValueError: if spec has unknown format
+	:returns imported broker:
+	"""
+	import_spec = object_spec.split(":")
+	if len(import_spec) != 2:
+		raise ValueError("you should provide object path in `module:variable` format.")
+	with add_cwd_in_path():
+		if app_dir:
+			sys.path.insert(0, app_dir)
+		module = import_module(import_spec[0])
+	return getattr(module, import_spec[1])
+
+
+def import_from_modules(modules: list[str]) -> None:
+	"""
+	import all modules from modules variable
+
+	:param modules: list of modules.
+	"""
+	for module in modules:
+		try:
+			logger.info(f"importing tasks from module {module}")
+			with add_cwd_in_path():
+				import_module(module)
+		except ImportError as err:
+			logger.warning(f"cannot import {module}. Cause:")
+			logger.exception(err)
+
+
+def import_tasks(modules: list[str], pattern: str | Sequence[str], fs_discover: bool) -> None:
+	"""
+	import tasks modules
+
+	this function is used to
+	import all tasks from modules
+
+	:param modules: list of modules to import
+	:param pattern: pattern of a file if fs_discover is true.
+	:param fs_discover: if true it will try to import modules
+	    from filesystem.
+	"""
+	if fs_discover:
+		if isinstance(pattern, str):
+			pattern = (pattern,)
+		discovered_modules = set()
+		for glob_pattern in pattern:
+			for path in Path().glob(glob_pattern):
+				if path.is_file():
+					if path.suffix in (".py", ".pyc", ".pyd", ".so"):
+						# remove all suffixes
+						prefix = path.name.partition(".")[0]
+						discovered_modules.add(
+							str(path.with_name(prefix)).replace(os.path.sep, ".")
+						)
+					# ignore other files
+				else:
+					discovered_modules.add(str(path).replace(os.path.sep, "."))
+
+		modules.extend(list(discovered_modules))
+	import_from_modules(modules)
+
+
+def object_fqn(obj: object) -> str:
+	if hasattr(obj, "__name__"):
+		return f"{obj.__module__}.{obj.__name__}"
+	return f"{obj.__module__}.{obj.__class__.__name__}"
+
+
+def get_type_fqn(arg: Any) -> str | None:
+	_resolved_module = ""
+	try:
+		_resolved_module = arg.__module__
+	except AttributeError:
+		if arg.__class__.__name__ in __builtins__:
+			_resolved_module = "builtins"
+
+	if _resolved_module == "":
+		return None
+
+	return _resolved_module + ":" + arg.__class__.__name__
+
+
+def get_type_from_fqn(_result: str | bytes | None) -> Any:
+	_imported_type = None
+	if _result is None:
+		return _imported_type
+
+	_decoded_result = _result.decode() if isinstance(_result, bytes) else _result
+	try:
+		_imported_type = import_object(_decoded_result)
+	except Exception as exc:
+		logger.warning("{}", exc)
+
+	return _imported_type

kuu-0.0.1/kuu/_types.py

@@ -0,0 +1,9 @@
+from typing import TYPE_CHECKING, Callable, Coroutine
+
+if TYPE_CHECKING:
+	from kuu.task import Task
+
+type _FnAsync[**P, R] = Callable[P, Coroutine[None, None, R]]
+type _Fn[**P, R] = Callable[P, R] | _FnAsync[P, R]
+type _FnSingle[P, R] = Callable[[P], R]
+type _Wrap[**P, R] = _FnSingle[_Fn[P, R], Task[P, R]]

kuu-0.0.1/kuu/app.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import inspect
+from datetime import datetime, timezone
+from typing import Any, overload
+
+from kuu._import import object_fqn
+from kuu._types import _Fn, _FnAsync, _Wrap
+from kuu.brokers.base import Broker
+from kuu.context import Context
+from kuu.events import Events
+from kuu.handle import TaskHandle
+from kuu.message import Message, Payload
+from kuu.middleware.base import Middleware, run_chain
+from kuu.registry import Registry
+from kuu.results.base import ResultBackend
+from kuu.serializers.base import Serializer
+from kuu.serializers.json import JSONSerializer
+from kuu.task import Task
+
+
+class Kuu:
+	def __init__(
+		self,
+		broker: Broker,
+		default_queue: str = "default",
+		middleware: list[Middleware] | None = None,
+		results: ResultBackend | None = None,
+		serializer: Serializer = JSONSerializer(),
+		result_ttl: float = 86400,
+		result_replay: bool = True,
+		result_store_errors: bool = True,
+	):
+		"""Initialize task manager
+
+		Args:
+				broker (Broker): broker implementation to handle tasks
+				default_queue (str, optional): default queue to communicate through. Defaults to "default".
+				middleware (list[Middleware] | None, optional): middlewares. Defaults to None.
+				results (ResultBackend | None, optional): result backend implementation. Defaults to None.
+				serializer (Serializer, optional): serializer implementation. Defaults to JSONSerializer().
+				result_ttl (float, optional): ttl for storing results. Defaults to 86400.
+				result_replay (bool, optional): Defaults to True.
+				result_store_errors (bool, optional): Defaults to True.
+		"""
+		self.broker = broker
+		self.results = results
+		self.serializer = serializer
+		self.default_queue = default_queue
+		self.middleware: list[Middleware] = list(middleware or [])
+		self.registry = Registry()
+		self.events = Events()
+		self.result_ttl = result_ttl
+		self.result_replay = result_replay
+		self.result_store_errors = result_store_errors
+
+	@overload
+	def task[**P, R](
+		self,
+		name: str | None = ...,
+		/,
+		queue: str | None = ...,
+		max_attempts: int = ...,
+		timeout: float | None = ...,
+		blocking: bool = ...,
+		**labels: Any,
+	) -> _Wrap[P, R]: ...
+
+	@overload
+	def task[**P, R](self, func: _Fn[P, R] = ..., /) -> Task[P, R]: ...
+
+	def task[**P, R](
+		self,
+		name_or_func: str | _Fn[P, R] | None = None,
+		/,
+		queue: str | None = None,
+		max_attempts: int = 5,
+		timeout: float | None = None,
+		blocking: bool = False,
+		**labels: Any,
+	) -> _Wrap[P, R] | Task[P, R]:
+
+		def _get_wrap(
+			_name: str | None = None,
+		):
+			def _wrap(func: _Fn[P, R]) -> Task[P, R]:
+				t: Task[P, R] = Task(
+					manager=self,
+					original_func=func,
+					task_name=_name or object_fqn(func),
+					task_queue=queue or self.default_queue,
+					task_labels=labels,
+					max_attempts=max_attempts,
+					timeout=timeout,
+					blocking=blocking,
+				)
+				self.registry.add(t)
+				return t
+
+			return _wrap
+
+		if inspect.isfunction(name_or_func):
+			func = name_or_func
+			return _get_wrap()(func)
+
+		name = name_or_func
+		return _get_wrap(name)
+
+	def _build_message(
+		self,
+		task_name: str,
+		task: Task | None,
+		args: Payload,
+		queue: str | None,
+		not_before: datetime | None,
+		headers: dict[str, str] | None,
+		max_attempts: int | None,
+	) -> Message:
+		return Message(
+			task=task_name,
+			queue=queue or (task.task_queue if task else self.default_queue),
+			payload=args,
+			headers=headers or {},
+			max_attempts=(
+				max_attempts if max_attempts is not None else (task.max_attempts if task else 5)
+			),
+			not_before=not_before,
+		)
+
+	async def _dispatch(
+		self,
+		msg: Message,
+		task: Task[Any, Any] | None,
+		not_before: datetime | None,
+	) -> None:
+		ctx = Context(app=self, message=msg, phase="enqueue", task=task)
+
+		async def _terminal(_c: Context) -> None:
+			if not_before is not None and not_before > datetime.now(timezone.utc):
+				await self.broker.schedule(msg, not_before)
+			else:
+				await self.broker.enqueue(msg)
+			await self.events.task_enqueued.send(msg)
+
+		await run_chain(ctx, self.middleware, _terminal)
+
+	def _enqueue_task[**P, Res](
+		self,
+		task: Task[P, Res],
+		queue: str | None = None,
+		not_before: datetime | None = None,
+		headers: dict[str, str] | None = None,
+		max_attempts: int | None = None,
+	) -> _FnAsync[P, TaskHandle[Res]]:
+		async def _(*args: P.args, **kwargs: P.kwargs) -> TaskHandle[Res]:
+			payload = Payload(args=args, kwargs=kwargs)
+			msg = self._build_message(
+				task.task_name,
+				task,
+				payload,
+				queue=queue,
+				not_before=not_before,
+				headers=headers,
+				max_attempts=max_attempts,
+			)
+			await self._dispatch(msg, task, not_before)
+			return TaskHandle[Res](message=msg, app=self)
+
+		return _
+
+	async def enqueue_by_name(
+		self,
+		task: str,
+		args: Payload = Payload(),
+		queue: str | None = None,
+		not_before: datetime | None = None,
+		headers: dict[str, str] | None = None,
+		max_attempts: int | None = None,
+	) -> TaskHandle[Any]:
+		t = self.registry.get(task)
+		msg = self._build_message(
+			task,
+			t,
+			args,
+			queue=queue,
+			not_before=not_before,
+			headers=headers,
+			max_attempts=max_attempts,
+		)
+		await self._dispatch(msg, t, not_before)
+		return TaskHandle[Any](message=msg, app=self)

kuu-0.0.1/kuu/brokers/__init__.py

@@ -0,0 +1,3 @@
+from .base import Broker, Delivery
+
+__all__ = ["Broker", "Delivery"]

kuu-0.0.1/kuu/brokers/base.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Protocol
+
+from ..message import Message
+
+
+@dataclass
+class Delivery[Receipt]:
+	message: Message
+	receipt: Receipt
+	queue: str
+
+
+class Broker(Protocol):
+	async def connect(self) -> None: ...
+	async def close(self) -> None: ...
+
+	async def declare(self, queue: str) -> None: ...
+
+	async def enqueue(self, msg: Message) -> None: ...
+	async def schedule(self, msg: Message, not_before: datetime) -> None: ...
+
+	def consume(self, queues: list[str], prefetch: int) -> AsyncIterator[Delivery]: ...
+
+	async def ack(self, delivery: Delivery) -> None: ...
+	async def nack(
+		self, delivery: Delivery, requeue: bool = True, delay: float | None = None
+	) -> None: ...

kuu-0.0.1/kuu/brokers/memory.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import heapq
+import itertools
+from collections.abc import AsyncIterator
+from datetime import datetime, timezone
+from typing import NamedTuple
+
+import anyio
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+
+from ..message import Message
+from .base import Broker, Delivery
+
+
+class MemoryReceipt(NamedTuple):
+	queue: str
+	seq: int
+
+
+class _Q(NamedTuple):
+	send: MemoryObjectSendStream[Delivery[MemoryReceipt]]
+	recv: MemoryObjectReceiveStream[Delivery[MemoryReceipt]]
+
+
+class MemoryBroker(Broker):
+	def __init__(self, buffer: int = 1024, pump_interval: float = 0.05):
+		self.buffer = buffer
+		self.pump_interval = pump_interval
+		self._queues: dict[str, _Q] = {}
+		# (run_at_ts, seq, queue, message)
+		self._scheduled: list[tuple[float, int, str, Message]] = []
+		self._sched_lock = anyio.Lock()
+		self._sched_event = anyio.Event()
+		self._seq = itertools.count()
+		self._pending: dict[tuple[str, int], Message] = {}
+
+	async def connect(self) -> None:
+		return None
+
+	async def close(self) -> None:
+		for q in self._queues.values():
+			await q.send.aclose()
+			await q.recv.aclose()
+		self._queues.clear()
+		self._scheduled.clear()
+		self._pending.clear()
+
+	async def declare(self, queue: str) -> None:
+		if queue in self._queues:
+			return
+		send, recv = anyio.create_memory_object_stream[Delivery[MemoryReceipt]](self.buffer)
+		self._queues[queue] = _Q(send, recv)
+
+	def _now(self) -> float:
+		return datetime.now(timezone.utc).timestamp()
+
+	def _ts(self, dt: datetime) -> float:
+		return dt.replace(tzinfo=timezone.utc).timestamp() if dt.tzinfo is None else dt.timestamp()
+
+	async def _push(self, queue: str, msg: Message) -> None:
+		await self.declare(queue)
+		seq = next(self._seq)
+		receipt = MemoryReceipt(queue=queue, seq=seq)
+		self._pending[(queue, seq)] = msg
+		await self._queues[queue].send.send(Delivery(message=msg, receipt=receipt, queue=queue))
+
+	async def enqueue(self, msg: Message) -> None:
+		await self._push(msg.queue, msg)
+
+	async def schedule(self, msg: Message, not_before: datetime) -> None:
+		await self.declare(msg.queue)
+		async with self._sched_lock:
+			heapq.heappush(self._scheduled, (self._ts(not_before), next(self._seq), msg.queue, msg))
+			self._sched_event.set()
+			self._sched_event = anyio.Event()
+
+	async def _pump_scheduled(self) -> None:
+		while True:
+			async with self._sched_lock:
+				now = self._now()
+				due: list[tuple[str, Message]] = []
+				while self._scheduled and self._scheduled[0][0] <= now:
+					_, _, q, m = heapq.heappop(self._scheduled)
+					due.append((q, m))
+			for q, m in due:
+				await self._push(q, m)
+			await anyio.sleep(self.pump_interval)
+
+	async def consume(
+		self, queues: list[str], prefetch: int
+	) -> AsyncIterator[Delivery[MemoryReceipt]]:
+		del prefetch  # in-memory has its own buffer
+		for q in queues:
+			await self.declare(q)
+
+		send, recv = anyio.create_memory_object_stream[Delivery[MemoryReceipt]](self.buffer)
+
+		async def _forward(q: str) -> None:
+			async for delivery in self._queues[q].recv.clone():
+				await send.send(delivery)
+
+		async with anyio.create_task_group() as tg:
+			tg.start_soon(self._pump_scheduled)
+			for q in queues:
+				tg.start_soon(_forward, q)
+			try:
+				async with recv:
+					async for delivery in recv:
+						yield delivery
+			finally:
+				tg.cancel_scope.cancel()
+
+	async def ack(self, delivery: Delivery) -> None:
+		match delivery.receipt:
+			case MemoryReceipt():
+				pass
+			case _:
+				return
+		self._pending.pop((delivery.receipt.queue, delivery.receipt.seq), None)
+
+	async def nack(
+		self,
+		delivery: Delivery,
+		requeue: bool = True,
+		delay: float | None = None,
+	) -> None:
+		match delivery.receipt:
+			case MemoryReceipt():
+				pass
+			case _:
+				return
+		key = (delivery.receipt.queue, delivery.receipt.seq)
+		self._pending.pop(key, None)
+		if not requeue:
+			return
+		msg = delivery.message.model_copy(update={"attempt": delivery.message.attempt + 1})
+		if delay and delay > 0:
+			when = datetime.fromtimestamp(self._now() + delay, tz=timezone.utc)
+			await self.schedule(msg, when)
+		else:
+			await self._push(delivery.receipt.queue, msg)