# SPDX-FileCopyrightText: 2026 Dennis Fink # # SPDX-License-Identifier: BSD-3-Clause from collections.abc import Callable from functools import wraps from typing import Any, ParamSpec, TypeVar from urllib.parse import urlsplit, urlunsplit from flask import Flask from huey import Huey as BaseHuey from huey.api import TaskWrapper P = ParamSpec("P") R = TypeVar("R") class Huey: """ Provide Flask integration for a Huey instance. The extension initializes a Huey backend from the Flask configuration and wraps tasks so that they execute within an application context. """ def __init__(self, app: Flask | None = None) -> None: """ Initialize the Huey extension. :param app: Flask application to initialize immediately, if provided. """ self.app: Flask | None = None self._huey: BaseHuey | None = None if app is not None: self.init_app(app) def init_app(self, app: Flask) -> None: """ Initialize Huey for a Flask application. The backend and its storage options are derived from ``HUEY_URL`` and the resulting extension is registered with the application. :param app: Flask application to initialize. :raises RuntimeError: If the configured Huey backend URL is invalid or uses an unsupported scheme. """ config: dict[str, Any] = { "name": app.import_name, "results": True, "store_none": False, "utc": True, "immediate": app.config.get("TESTING", False), **app.config.get("HUEY", {}), } url = app.config.get("HUEY_URL", config.pop("url", "memory://")) huey_class, storage_kwargs = self.backend_from_url(url) self.app = app self._huey = huey_class(**config, **storage_kwargs) app.extensions["huey"] = self @property def huey(self) -> BaseHuey: """ Return the initialized Huey instance. :return: Configured Huey backend instance. :raises RuntimeError: If the extension has not been initialized. """ if self._huey is None: raise RuntimeError( "Huey has not been initialized. " "Call huey.init_app(app) before importing tasks." ) return self._huey def task( self, *task_args: Any, **task_kwargs: Any ) -> Callable[[Callable[P, object]], TaskWrapper]: """ Create a Huey task that runs within the Flask application context. :param task_args: Positional arguments forwarded to Huey's task decorator. :param task_kwargs: Keyword arguments forwarded to Huey's task decorator. :return: Decorator that registers the wrapped function as a Huey task. """ def decorator(func: Callable[P, R]) -> TaskWrapper: @wraps(func) def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: if self.app is None: raise RuntimeError("Flask app is not available.") with self.app.app_context(): return func(*args, **kwargs) return self.huey.task(*task_args, **task_kwargs)(wrapper) return decorator def periodic_task( self, *task_args: Any, **task_kwargs: Any ) -> Callable[[Callable[P, object]], TaskWrapper]: """ Create a periodic Huey task that runs within the Flask application context. :param task_args: Positional arguments forwarded to Huey's periodic task decorator. :param task_kwargs: Keyword arguments forwarded to Huey's periodic task decorator. :return: Decorator that registers the wrapped function as a periodic Huey task. """ def decorator(func: Callable[P, R]) -> TaskWrapper: @wraps(func) def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: if self.app is None: raise RuntimeError("Flask app is not available.") with self.app.app_context(): return func(*args, **kwargs) return self.huey.periodic_task(*task_args, **task_kwargs)(wrapper) return decorator def __getattr__(self, name: str) -> Any: """ Forward an unknown attribute to the underlying Huey instance. :param name: Name of the attribute to retrieve. :return: Attribute from the initialized Huey instance. :raises RuntimeError: If the extension has not been initialized. """ return getattr(self.huey, name) @staticmethod def backend_from_url(url: str) -> tuple[type[BaseHuey], dict[str, Any]]: """ Determine the Huey backend and storage options from a URL. :param url: Huey backend URL. :return: Huey backend class and keyword arguments for its storage backend. :raises RuntimeError: If the URL is malformed for the selected backend or uses an unsupported scheme. """ parsed = urlsplit(url) scheme = parsed.scheme.lower() if scheme.startswith(("redis", "rediss")): fixed_url = urlunsplit(parsed._replace(scheme=scheme.split("+", 1)[0])) if scheme.endswith("priority+expire"): from huey import PriorityRedisExpireHuey return PriorityRedisExpireHuey, {"url": fixed_url} elif scheme.endswith("priority"): from huey import PriorityRedisHuey return PriorityRedisHuey, {"url": fixed_url} elif scheme.endswith("expire"): from huey import RedisExpireHuey return RedisExpireHuey, {"url": fixed_url} else: from huey import RedisHuey return RedisHuey, {"url": url} elif scheme == "sqlite": from huey import SqliteHuey prefix = "sqlite:///" if not url.startswith(prefix): raise RuntimeError( "SQLite Huey URLs must look like sqlite:///var/huey.db" ) filename = url.removeprefix(prefix) if not filename: raise RuntimeError("SQLite Huey URL must include a database path.") return SqliteHuey, {"filename": filename} elif scheme == "file": from huey import FileHuey prefix = "file:///" if not url.startswith(prefix): raise RuntimeError( "File Huey URLs must look like file:///var/huey-queue" ) path = url.removeprefix(prefix) if not path: raise RuntimeError("File Huey URL must include a directory path.") return FileHuey, {"path": path} elif scheme in {"postgres", "postgresql"}: from huey import PostgresHuey return PostgresHuey, {"dsn": url} elif scheme == "memory": from huey import MemoryHuey return MemoryHuey, {} elif scheme == "blackhole": from huey import BlackHoleHuey return BlackHoleHuey, {} raise RuntimeError(f"Unsupported HUEY_URL scheme: {scheme!r}")