singleton.py

"""Singleton decorator supporting classes, functions, and async/sync operations."""

import asyncio
import functools
import threading
from collections.abc import Awaitable, Callable
from typing import Any, ClassVar, ParamSpec, TypeVar

P = ParamSpec("P")
T = TypeVar("T")
R = TypeVar("R")


class _SingletonMeta(type):
    """Metaclass for thread-safe singleton classes."""

    _instances: ClassVar[dict[type[Any], Any]] = {}
    _locks: ClassVar[dict[type[Any], threading.Lock]] = {}
    _lock: ClassVar[threading.Lock] = threading.Lock()

    def __call__(cls: type[Any], *args: Any, **kwargs: Any) -> Any:
        if cls not in cls._instances:
            with cls._lock:
                if cls not in cls._locks:
                    cls._locks[cls] = threading.Lock()
                if cls not in cls._instances:
                    cls._instances[cls] = super().__call__(*args, **kwargs)
        return cls._instances[cls]


class _FunctionSingleton:
    """Wrapper for singleton function calls."""

    def __init__(self, func: Callable[..., Any]) -> None:
        self.func = func
        functools.update_wrapper(self, func)
        self._result: Any | None = None
        self._lock = threading.Lock()
        self._called = False

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if not self._called:
            with self._lock:
                if not self._called:
                    self._result = self.func(*args, **kwargs)
                    self._called = True
        if self._result is None:
            msg = "Function result is None"
            raise RuntimeError(msg)
        return self._result


class _AsyncFunctionSingleton:
    """Wrapper for singleton async function calls."""

    def __init__(self, func: Callable[..., Awaitable[Any]]) -> None:
        self.func = func
        functools.update_wrapper(self, func)
        self._result: Any | None = None
        self._lock: asyncio.Lock | None = None
        self._called = False
        self._init_lock = threading.Lock()

    async def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if not self._called:
            with self._init_lock:
                if self._lock is None:
                    self._lock = asyncio.Lock()
            if self._lock is None:
                msg = "Lock initialization failed"
                raise RuntimeError(msg)
            async with self._lock:
                if not self._called:
                    self._result = await self.func(*args, **kwargs)
                    self._called = True
        if self._result is None:
            msg = "Function result is None"
            raise RuntimeError(msg)
        return self._result


def singleton(
    obj: type[T] | Callable[P, R] | Callable[P, Awaitable[R]] | None = None,
    *,
    async_safe: bool = False,
) -> Any:
    """Singleton decorator supporting classes, functions, and async/sync operations.

    Usage:
        # For classes:
        @singleton
        class MyClass:
            pass

        # For functions:
        @singleton
        def my_function():
            return expensive_operation()

        # For async functions:
        @singleton
        async def my_async_function():
            return await expensive_async_operation()

    Args:
        obj: The class or function to make singleton.
        async_safe: If True, uses async locks for thread-safe async operations.
            Currently only affects function behavior, classes use thread-safe locks.

    Returns:
        Singleton instance of the class or cached result of the function.

    """
    if obj is None:
        # Decorator factory mode: @singleton(async_safe=True)
        def decorator(target: Any) -> Any:
            return singleton(target, async_safe=async_safe)

        return decorator

    # Direct decorator mode: @singleton
    if isinstance(obj, type):
        # It's a class - use metaclass for thread-safe singleton
        return type(
            obj.__name__,
            (obj,),
            {"__metaclass__": _SingletonMeta},
        )

    # It's a function
    if asyncio.iscoroutinefunction(obj):
        # Async function
        return _AsyncFunctionSingleton(obj)

    # Sync function
    return _FunctionSingleton(obj)