wool 0.1rc20__py3-none-any.whl

wool/__init__.py

@@ -0,0 +1,122 @@
+from contextvars import ContextVar
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version
+from typing import Final
+
+from tblib import pickling_support
+
+from wool._context import AppContext
+from wool._resource_pool import ResourcePool
+from wool._work import WoolTask
+from wool._work import WoolTaskEvent
+from wool._work import WoolTaskEventCallback
+from wool._work import WoolTaskEventType
+from wool._work import WoolTaskException
+from wool._work import current_task as wool_current_task
+from wool._work import work
+from wool.core.discovery.base import Discovery
+from wool.core.discovery.base import DiscoveryEvent
+from wool.core.discovery.base import DiscoveryEventType
+from wool.core.discovery.base import DiscoveryLike
+from wool.core.discovery.base import DiscoveryPublisherLike
+from wool.core.discovery.base import DiscoverySubscriberLike
+from wool.core.discovery.base import PredicateFunction
+from wool.core.discovery.base import WorkerInfo
+from wool.core.discovery.lan import LanDiscovery
+from wool.core.discovery.local import LocalDiscovery
+from wool.core.loadbalancer.base import ConnectionResourceFactory
+from wool.core.loadbalancer.base import LoadBalancerContext
+from wool.core.loadbalancer.base import LoadBalancerLike
+from wool.core.loadbalancer.base import NoWorkersAvailable
+from wool.core.loadbalancer.roundrobin import RoundRobinLoadBalancer
+from wool.core.typing import Factory
+from wool.core.worker.base import Worker
+from wool.core.worker.base import WorkerFactory
+from wool.core.worker.base import WorkerLike
+from wool.core.worker.connection import RpcError
+from wool.core.worker.connection import TransientRpcError
+from wool.core.worker.connection import UnexpectedResponse
+from wool.core.worker.connection import WorkerConnection
+from wool.core.worker.local import LocalWorker
+from wool.core.worker.pool import WorkerPool
+from wool.core.worker.proxy import WorkerProxy
+from wool.core.worker.service import WorkerService
+
+pickling_support.install()
+
+# Alias for backwards compatibility
+routine = work
+
+try:
+    __version__ = version("wool")
+except PackageNotFoundError:
+    __version__ = "unknown"
+
+__proxy__: Final[ContextVar[WorkerProxy | None]] = ContextVar("__proxy__", default=None)
+
+__proxy_pool__: Final[ContextVar[ResourcePool[WorkerProxy] | None]] = ContextVar(
+    "__proxy_pool__", default=None
+)
+
+__all__ = [
+    # Connection
+    "RpcError",
+    "TransientRpcError",
+    "UnexpectedResponse",
+    "WorkerConnection",
+    # Context
+    "AppContext",
+    # Load balancing
+    "ConnectionResourceFactory",
+    "LoadBalancerContext",
+    "LoadBalancerLike",
+    "NoWorkersAvailable",
+    "RoundRobinLoadBalancer",
+    # Work
+    "WoolTask",
+    "WoolTaskEvent",
+    "WoolTaskEventCallback",
+    "WoolTaskEventType",
+    "WoolTaskException",
+    "routine",
+    "work",
+    "wool_current_task",
+    # Workers
+    "LocalWorker",
+    "Worker",
+    "WorkerFactory",
+    "WorkerLike",
+    "WorkerPool",
+    "WorkerProxy",
+    "WorkerService",
+    # Discovery
+    "Discovery",
+    "DiscoveryEvent",
+    "DiscoveryEventType",
+    "DiscoveryLike",
+    "DiscoveryPublisherLike",
+    "DiscoverySubscriberLike",
+    "LanDiscovery",
+    "LocalDiscovery",
+    "PredicateFunction",
+    "WorkerInfo",
+    # Typing
+    "Factory",
+]
+
+for symbol in __all__:
+    attribute = globals().get(symbol)
+    try:
+        if attribute and "wool" in attribute.__module__.split("."):
+            # Set the module to reflect imports of the symbol
+            attribute.__module__ = __name__
+    except AttributeError:
+        continue
+
+# for plugin in entry_points(group="wool_cli_plugins"):
+#     try:
+#         plugin.load()
+#         logging.info(f"Loaded CLI plugin {plugin.name}")
+#     except Exception as e:
+#         logging.error(f"Failed to load CLI plugin {plugin.name}: {e}")
+#         raise

wool/_context.py

@@ -0,0 +1,29 @@
+from contextvars import ContextVar
+from contextvars import Token
+from typing import Final
+
+from wool._undefined import Undefined
+from wool._undefined import UndefinedType
+
+dispatch_timeout: Final[ContextVar[float | None]] = ContextVar(
+    "_dispatch_timeout", default=None
+)
+
+
+# public
+class AppContext:
+    _dispatch_timeout: float | None | UndefinedType
+    _dispatch_timeout_token: Token | UndefinedType
+
+    def __init__(self, *, dispatch_timeout: float | None | UndefinedType = Undefined):
+        self._dispatch_timeout = dispatch_timeout
+
+    def __enter__(self):
+        if self._dispatch_timeout is not Undefined:
+            self._dispatch_timeout_token = dispatch_timeout.set(self._dispatch_timeout)
+        else:
+            self._dispatch_timeout_token = Undefined
+
+    def __exit__(self, *_):
+        if self._dispatch_timeout_token is not Undefined:
+            dispatch_timeout.reset(self._dispatch_timeout_token)

wool/_protobuf/worker.py

@@ -0,0 +1,26 @@
+try:
+    from wool._protobuf.worker_pb2 import Ack
+    from wool._protobuf.worker_pb2 import Nack
+    from wool._protobuf.worker_pb2 import Response
+    from wool._protobuf.worker_pb2 import StopRequest
+    from wool._protobuf.worker_pb2 import Void
+    from wool._protobuf.worker_pb2 import WorkerInfo
+    from wool._protobuf.worker_pb2_grpc import WorkerServicer
+    from wool._protobuf.worker_pb2_grpc import WorkerStub
+    from wool._protobuf.worker_pb2_grpc import add_WorkerServicer_to_server
+except ImportError as e:
+    from wool._protobuf.exception import ProtobufImportError
+
+    raise ProtobufImportError(e) from e
+
+__all__ = [
+    "Ack",
+    "Nack",
+    "Response",
+    "StopRequest",
+    "Void",
+    "WorkerInfo",
+    "WorkerServicer",
+    "WorkerStub",
+    "add_WorkerServicer_to_server",
+]

wool/_resource_pool.py

@@ -0,0 +1,376 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import Any
+from typing import Awaitable
+from typing import Callable
+from typing import Final
+from typing import Generic
+from typing import TypeVar
+from typing import cast
+
+T = TypeVar("T")
+
+
+SENTINEL: Final = object()
+
+
+class Resource(Generic[T]):
+    """
+    A single-use async context manager for resource acquisition.
+
+    This class can only be used once as an async context manager. After
+    acquisition, it cannot be reacquired, and after release, it cannot be
+    released again.
+
+    :param pool:
+        The :class:`ResourcePool` this resource belongs to.
+    :param key:
+        The cache key for this resource.
+    """
+
+    def __init__(self, pool: ResourcePool[T], key):
+        self._pool = pool
+        self._key = key
+        self._resource = None
+        self._acquired = False
+        self._released = False
+
+    async def __aenter__(self) -> T:
+        """
+        Context manager entry - acquire resource.
+
+        :returns:
+            The cached resource object.
+        :raises RuntimeError:
+            If called on a resource that was previously acquired.
+        """
+        if self._acquired:
+            raise RuntimeError(
+                "Cannot re-acquire a resource that has already been acquired"
+            )
+
+        self._acquired = True
+        try:
+            self._resource = await self._pool.acquire(self._key)
+            return cast(T, self._resource)
+        except Exception:
+            self._acquired = False
+            raise
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """
+        Context manager exit - release resource.
+
+        :param exc_type:
+            Exception type if an exception occurred, None otherwise.
+        :param exc_val:
+            Exception value if an exception occurred, None otherwise.
+        :param exc_tb:
+            Exception traceback if an exception occurred, None otherwise.
+        """
+        await self._release()
+
+    async def _release(self):
+        """
+        Release the resource.
+
+        :raises RuntimeError:
+            If attempting to release a resource that was not acquired or
+            already released.
+        """
+        if not self._acquired:
+            raise RuntimeError("Cannot release a resource that was not acquired")
+        if self._released:
+            raise RuntimeError(
+                "Cannot release a resource that has already been released"
+            )
+
+        self._released = True
+        if self._resource:
+            await self._pool.release(self._key)
+
+
+class ResourcePool(Generic[T]):
+    """
+    An asynchronous reference-counted cache with TTL-based cleanup.
+
+    Objects are created on-demand via a factory function (sync or async) and
+    automatically cleaned up after all references are released and the TTL
+    expires.
+
+    :param factory:
+        Function to create new objects (sync or async).
+    :param finalizer:
+        Optional cleanup function (sync or async).
+    :param ttl:
+        Time-to-live in seconds after last reference is released.
+    """
+
+    @dataclass
+    class CacheEntry:
+        """
+        Internal cache entry tracking an object and its metadata.
+
+        :param obj:
+            The cached object.
+        :param reference_count:
+            Number of active references to this object.
+        :param cleanup:
+            Optional cleanup task scheduled when reference count reaches zero.
+        """
+
+        obj: Any
+        reference_count: int
+        cleanup: asyncio.Task | None = None
+
+    @dataclass
+    class Stats:
+        """
+        Statistics about the current state of the resource pool.
+
+        :param total_entries:
+            Total number of cached entries.
+        :param referenced_entries:
+            Number of entries currently being referenced (reference_count > 0).
+        :param pending_cleanup:
+            Number of cleanup tasks currently pending execution.
+        """
+
+        total_entries: int
+        referenced_entries: int
+        pending_cleanup: int
+
+    def __init__(
+        self,
+        factory: Callable[[Any], T | Awaitable[T]],
+        *,
+        finalizer: Callable[[T], None | Awaitable[None]] | None = None,
+        ttl: float = 0,
+    ):
+        self._factory = factory
+        self._finalizer = finalizer
+        self._ttl = ttl
+        self._cache: dict[Any, ResourcePool.CacheEntry] = {}
+        self._lock = asyncio.Lock()
+
+    async def __aenter__(self):
+        """Async context manager entry.
+
+        :returns:
+            The ResourcePool instance itself.
+        """
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit - cleanup all resources.
+
+        :param exc_type:
+            Exception type if an exception occurred, None otherwise.
+        :param exc_val:
+            Exception value if an exception occurred, None otherwise.
+        :param exc_tb:
+            Exception traceback if an exception occurred, None otherwise.
+        """
+        await self.clear()
+
+    @property
+    def stats(self) -> Stats:
+        """
+        Return cache statistics.
+
+        .. note::
+            This is synchronous for convenience, but should only be called
+            when not concurrently modifying the cache.
+
+        :returns:
+            :class:`ResourcePool.Stats` containing current statistics.
+        """
+        pending_cleanup = sum(
+            1 for c in self.pending_cleanup.values() if c is not None and not c.done()
+        )
+        return self.Stats(
+            total_entries=len(self._cache),
+            referenced_entries=sum(
+                1 for e in self._cache.values() if e.reference_count > 0
+            ),
+            pending_cleanup=pending_cleanup,
+        )
+
+    @property
+    def pending_cleanup(self):
+        """Dictionary of cache keys with pending cleanup tasks.
+
+        :returns:
+            Dictionary mapping cache keys to their cleanup tasks.
+        """
+        return {
+            k: v.cleanup
+            for k, v in self._cache.items()
+            if v.cleanup is not None and not v.cleanup.done()
+        }
+
+    def get(self, key: Any) -> Resource[T]:
+        """
+        Get a resource acquisition that can be awaited or used as context
+        manager.
+
+        :param key:
+            The cache key.
+        :returns:
+            :class:`Resource` that can be awaited or used with 'async with'.
+        """
+        return Resource(self, key)
+
+    async def acquire(self, key: Any) -> T:
+        """
+        Internal acquire method - acquires a reference to the cached object.
+
+        Creates a new object via the factory if not cached. Increments
+        reference count and cancels any pending cleanup.
+
+        :param key:
+            The cache key.
+        :returns:
+            The cached or newly created object.
+        """
+        async with self._lock:
+            if key in self._cache:
+                entry = self._cache[key]
+                entry.reference_count += 1
+
+                # Cancel pending cleanup task if it exists
+                if entry.cleanup is not None and not entry.cleanup.done():
+                    entry.cleanup.cancel()
+                    try:
+                        await entry.cleanup
+                    except asyncio.CancelledError:
+                        pass
+                    entry.cleanup = None
+
+                return entry.obj
+            else:
+                # Cache miss - create new object
+                obj = await self._await(self._factory, key)
+                self._cache[key] = self.CacheEntry(obj=obj, reference_count=1)
+                return obj
+
+    async def release(self, key: Any) -> None:
+        """
+        Release a reference to the cached object.
+
+        Decrements reference count. If count reaches 0, schedules cleanup
+        after TTL expires (if TTL > 0).
+
+        :param key:
+            The cache key.
+        :raises KeyError:
+            If key not in cache.
+        """
+        async with self._lock:
+            if key not in self._cache:
+                raise KeyError(f"Key '{key}' not found in cache")
+            entry = self._cache[key]
+
+            if entry.reference_count <= 0:
+                raise ValueError(f"Reference count for key '{key}' is already 0")
+
+            entry.reference_count -= 1
+
+            if entry.reference_count <= 0:
+                if self._ttl > 0:
+                    # Schedule cleanup after TTL
+                    entry.cleanup = asyncio.create_task(self._schedule_cleanup(key))
+                else:
+                    # Immediate cleanup
+                    await self._cleanup(key)
+
+    async def clear(self, key=SENTINEL) -> None:
+        """Clear cache entries and cancel pending cleanups.
+
+        :param key:
+            Specific key to clear, or SENTINEL to clear all entries.
+        """
+        async with self._lock:
+            # Clean up all entries
+            if key is SENTINEL:
+                keys = list(self._cache.keys())
+            else:
+                keys = [key]
+            for key in keys:
+                await self._cleanup(key)
+
+    async def _schedule_cleanup(self, key: Any) -> None:
+        """
+        Schedule cleanup after TTL delay.
+
+        Only cleans up if the reference count is still 0 when TTL expires.
+
+        :param key:
+            The cache key to schedule cleanup for.
+        """
+        try:
+            await asyncio.sleep(self._ttl)
+
+            async with self._lock:
+                # Double-check conditions - reference might have been re-acquired
+                if key in self._cache:
+                    entry = self._cache[key]
+                    if entry.reference_count == 0:
+                        await self._cleanup(key)
+
+        except asyncio.CancelledError:
+            # Cleanup was cancelled due to new reference - this is expected
+            pass
+
+    async def _cleanup(self, key: Any) -> None:
+        """
+        Remove entry from cache and call finalizer.
+
+        .. warning::
+            Must be called while holding the lock.
+
+        :param key:
+            The cache key to cleanup.
+        """
+        entry = self._cache[key]
+        try:
+            # Cancel cleanup task if running
+            if entry.cleanup is not None and not entry.cleanup.done():
+                entry.cleanup.cancel()
+                try:
+                    await entry.cleanup
+                except asyncio.CancelledError:
+                    pass
+        finally:
+            # Call finalizer
+            if self._finalizer:
+                try:
+                    await self._await(self._finalizer, entry.obj)
+                except Exception:
+                    pass
+            del self._cache[key]
+
+    async def _await(self, func: Callable, *args) -> Any:
+        """
+        Call a function that might be sync or async.
+
+        If the function is a coroutine function, await it. Otherwise, call it
+        synchronously. If the result is a coroutine, await that as well.
+
+        :param func:
+            The function to call.
+        :param args:
+            Arguments to pass to the function.
+        :returns:
+            The result of the function call.
+        """
+        if asyncio.iscoroutinefunction(func):
+            return await func(*args)
+        else:
+            result = func(*args)
+            # Check if the result is a coroutine and await it if so
+            if asyncio.iscoroutine(result):
+                return await result
+            return result

wool/_typing.py

@@ -0,0 +1,7 @@
+from typing import Callable
+from typing import TypeVar
+
+F = TypeVar("F", bound=Callable)
+W = TypeVar("W", bound=Callable)
+Decorator = Callable[[F], W]
+PassthroughDecorator = Callable[[F], F]

wool/_undefined.py

@@ -0,0 +1,11 @@
+from enum import Enum
+from typing import Final
+from typing import final
+
+
+@final
+class UndefinedType(Enum):
+    Undefined = "Undefined"
+
+
+Undefined: Final = UndefinedType.Undefined