wool 0.1rc9__py3-none-any.whl → 0.1rc10__py3-none-any.whl

wool/_manager.py

@@ -1,181 +0,0 @@
-from __future__ import annotations
-
-import logging
-from functools import partial
-from multiprocessing.managers import BaseManager
-from multiprocessing.managers import DictProxy
-from queue import Empty
-from threading import Event
-from threading import Lock
-from typing import TYPE_CHECKING
-from typing import Any
-from typing import Callable
-from typing import TypeVar
-from typing import overload
-from uuid import UUID
-from weakref import WeakValueDictionary
-
-from wool._event import TaskEvent
-from wool._future import Future
-from wool._queue import TaskQueue
-from wool._typing import PassthroughDecorator
-
-if TYPE_CHECKING:
-    from wool._task import Task
-
-
-C = TypeVar("C", bound=Callable[..., Any])
-
-_manager_registry = {}
-
-
-def _register(fn, /, **kwargs):
-    assert fn.__name__ not in _manager_registry, (
-        f"Function '{fn.__name__}' already registered"
-    )
-    _manager_registry[fn.__name__] = (fn, kwargs)
-    return fn
-
-
-@overload
-def register(
-    fn: C,
-    /,
-    *,
-    exposed: tuple[str, ...] | None = None,
-    proxytype: None = None,
-    method_to_typeid: None = None,
-) -> C: ...
-
-
-@overload
-def register(
-    fn: None = None,
-    /,
-    *,
-    exposed: tuple[str, ...] | None = None,
-    proxytype: type | None = None,
-    method_to_typeid: dict[str, str] | None = None,
-) -> PassthroughDecorator[C]: ...
-
-
-def register(
-    fn: C | None = None,
-    /,
-    *,
-    exposed: tuple[str, ...] | None = None,
-    proxytype: type | None = None,
-    method_to_typeid: dict[str, str] | None = None,
-) -> PassthroughDecorator[C] | C:
-    kwargs = {}
-    if exposed is not None:
-        kwargs["exposed"] = exposed
-    if proxytype is not None:
-        kwargs["proxytype"] = proxytype
-    if method_to_typeid is not None:
-        kwargs["method_to_typeid"] = method_to_typeid
-    if fn:
-        return _register(fn, **kwargs)
-    else:
-        return partial(_register, **kwargs)
-
-
-class FuturesProxy(DictProxy):
-    assert (method_to_typeid := getattr(DictProxy, "_method_to_typeid_"))
-    _method_to_typeid_ = {
-        **method_to_typeid,
-        "__getitem__": "proxify",
-        "setdefault": "proxify",
-    }
-
-
-_task_queue: TaskQueue[Task] = TaskQueue(100000, None)
-
-_task_queue_lock = Lock()
-
-_task_futures: WeakValueDictionary[UUID, Future] = WeakValueDictionary()
-
-
-@register
-def put(task: Task) -> Future:
-    try:
-        with queue_lock():
-            queue().put(task, block=False)
-            future = futures()[task.id] = Future()
-            TaskEvent("task-queued", task=task).emit()
-            return future
-    except Exception as e:
-        logging.exception(e)
-        raise
-
-
-@register
-def get() -> Task | Empty | None:
-    try:
-        return queue().get(block=False)
-    except Empty as e:
-        return e
-
-
-@register
-def proxify(value):
-    return value
-
-
-@register(proxytype=FuturesProxy)
-def futures():
-    global _task_futures
-    if not _task_futures:
-        _task_futures = WeakValueDictionary()
-    return _task_futures
-
-
-@register
-def queue() -> TaskQueue:
-    global _task_queue
-    if not _task_queue:
-        _task_queue = TaskQueue(1000, None)
-    return _task_queue
-
-
-@register
-def queue_lock() -> Lock:
-    global _task_queue_lock
-    if not _task_queue_lock:
-        _task_queue_lock = Lock()
-    return _task_queue_lock
-
-
-_stop_event = Event()
-_wait_event = Event()
-
-
-@register(exposed=("is_set", "set", "clear", "wait"))
-def stopping() -> Event:
-    return _stop_event
-
-
-@register(exposed=("is_set", "set", "clear", "wait"))
-def waiting() -> Event:
-    return _wait_event
-
-
-class ManagerMeta(type):
-    def __new__(mcs, name, bases, attrs):
-        cls = super().__new__(mcs, name, bases, attrs)
-        assert (register := getattr(cls, "register"))
-        for name, (fn, kwargs) in _manager_registry.items():
-            register(name, callable=fn, **kwargs)
-        return cls
-
-
-class Manager(BaseManager, metaclass=ManagerMeta):
-    if TYPE_CHECKING:
-        put = staticmethod(put)
-        get = staticmethod(get)
-        proxify = staticmethod(proxify)
-        futures = staticmethod(futures)
-        queue = staticmethod(queue)
-        queue_lock = staticmethod(queue_lock)
-        stopping = staticmethod(stopping)
-        waiting = staticmethod(waiting)

wool/_mempool/__init__.py

@@ -1,4 +0,0 @@
-from wool._mempool._mempool import MemoryPool
-from wool._mempool._service import MemoryPoolService
-
-__all__ = ["MemoryPool", "MemoryPoolService"]

wool/_mempool/_client.py

@@ -1,167 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-from typing import Final
-from typing import Protocol
-
-try:
-    from typing import Self
-except ImportError:
-    from typing_extensions import Self
-
-import grpc
-
-try:
-    from wool._protobuf.mempool import service_pb2 as pb
-    from wool._protobuf.mempool import service_pb2_grpc as rpc
-except ImportError as e:
-    from wool._protobuf import ProtobufImportError
-
-    raise ProtobufImportError(e) from e
-
-
-class EventHandler(Protocol):
-    @staticmethod
-    def __call__(client: MemoryPoolClient, response: pb.Event) -> None: ...
-
-
-class MemoryPoolClient:
-    _event_stream: Final[grpc.aio.UnaryStreamCall]
-    _listener: asyncio.Task
-    _session: pb.Session | None
-
-    def __init__(self, channel: grpc.aio.Channel, event_handler: EventHandler):
-        self._stub = rpc.MemoryPoolStub(channel)
-        self._event_stream = self._stub.session(pb.SessionRequest())
-        self._event_handler = event_handler
-        self._session = None
-
-    async def __aenter__(self) -> Self:
-        if self._session:
-            raise RuntimeError(
-                "Client session may not be re-entered. "
-                "Use 'async with MemoryPoolClient(...)' only once."
-            )
-        response = await self._event_stream.read()
-        assert isinstance(response, pb.SessionResponse)
-        self._session = response.session
-        self._listener = asyncio.create_task(self._listen())
-        return self
-
-    async def __aexit__(self, *_):
-        if not self._session:
-            raise RuntimeError(
-                "Client session has not been entered. "
-                "Use 'async with MemoryPoolClient(...)'."
-            )
-        assert self._listener
-        self._listener.cancel()
-        try:
-            await self._listener
-        except asyncio.CancelledError:
-            pass
-
-    def __del__(self):
-        try:
-            self._event_stream.cancel()
-        except Exception:
-            pass
-
-    @property
-    def id(self) -> str:
-        if not self._session:
-            raise RuntimeError(
-                "Client session has not been entered. "
-                "Use 'async with MemoryPoolClient(...)'."
-            )
-        return self._session.id
-
-    async def map(self, ref: str):
-        if not self._session:
-            raise RuntimeError(
-                "Client session has not been entered. "
-                "Use 'async with MemoryPoolClient(...)'."
-            )
-        request = pb.AcquireRequest(
-            session=self._session,
-            reference=pb.Reference(id=ref),
-        )
-        await self._stub.map(request)
-
-    async def get(self, ref: str) -> bytes:
-        if not self._session:
-            raise RuntimeError(
-                "Client session has not been entered. "
-                "Use 'async with MemoryPoolClient(...)'."
-            )
-        request = pb.GetRequest(
-            session=self._session,
-            reference=pb.Reference(id=ref),
-        )
-        response: pb.GetResponse = await self._stub.get(request)
-        return response.dump
-
-    async def put(self, dump: bytes, *, mutable: bool = False) -> str:
-        if not self._session:
-            raise RuntimeError(
-                "Client session has not been entered. "
-                "Use 'async with MemoryPoolClient(...)'."
-            )
-        request = pb.PutRequest(
-            session=self._session,
-            dump=dump,
-            mutable=mutable,
-        )
-        response: pb.PutResponse = await self._stub.put(request)
-        return response.reference.id
-
-    async def post(self, ref: str, dump: bytes) -> bool:
-        if not self._session:
-            raise RuntimeError(
-                "Client session has not been entered. "
-                "Use 'async with MemoryPoolClient(...)'."
-            )
-        request = pb.PostRequest(
-            session=self._session,
-            reference=pb.Reference(id=ref),
-            dump=dump,
-        )
-        response: pb.PostResponse = await self._stub.post(request)
-        return response.updated
-
-    async def acquire(self, ref: str):
-        if not self._session:
-            raise RuntimeError(
-                "Client session has not been entered. "
-                "Use 'async with MemoryPoolClient(...)'."
-            )
-        request = pb.AcquireRequest(
-            session=self._session,
-            reference=pb.Reference(id=ref),
-        )
-        await self._stub.acquire(request)
-
-    async def release(self, ref: str):
-        if not self._session:
-            raise RuntimeError(
-                "Client session has not been entered. "
-                "Use 'async with MemoryPoolClient(...)'."
-            )
-        request = pb.ReleaseRequest(
-            session=self._session,
-            reference=pb.Reference(id=ref),
-        )
-        await self._stub.release(request)
-
-    async def _listen(self):
-        assert self._event_stream
-        try:
-            while True:
-                if (response := await self._event_stream.read()) is None:
-                    break
-                assert isinstance(response, pb.SessionResponse), (
-                    f"Unexpected event type: {type(response)}"
-                )
-                self._event_handler(self, response.event)
-        except asyncio.CancelledError:
-            self._event_stream.cancel()

wool/_mempool/_mempool.py

@@ -1,311 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-import hashlib
-import os
-import pathlib
-import shutil
-from contextlib import asynccontextmanager
-from mmap import mmap
-from typing import BinaryIO
-
-try:
-    from typing import Self
-except ImportError:
-    from typing_extensions import Self
-
-import shortuuid
-
-from wool._mempool._metadata import MetadataMessage
-
-
-class SharedObject:
-    _id: str
-    _mempool: MemoryPool
-    _file: BinaryIO
-    _mmap: mmap
-    _size: int
-    _md5: bytes
-
-    def __init__(self, id: str, *, mempool: MemoryPool):
-        self._id = id
-        self._mempool = mempool
-        self._file = open(self._path / "dump", "r+b")
-        self._mmap = mmap(self._file.fileno(), 0)
-        self._size = self.metadata.size
-        self._md5 = self.metadata.md5
-
-    def __del__(self):
-        try:
-            self.close()
-        except Exception:
-            pass
-
-    @property
-    def id(self) -> str:
-        return self._id
-
-    @property
-    def metadata(self) -> SharedObjectMetadata:
-        return SharedObjectMetadata(self.id, mempool=self._mempool)
-
-    @property
-    def mmap(self) -> mmap:
-        return self._mmap
-
-    @property
-    def _path(self) -> pathlib.Path:
-        return pathlib.Path(self._mempool.path, self.id)
-
-    def close(self):
-        self.metadata.close()
-        self._mmap.close()
-        self._file.close()
-
-    def refresh(self) -> Self:
-        if self._size != self.metadata.size or self._md5 != self.metadata.md5:
-            self._mmap.close()
-            self._file.close()
-            self._file = open(self._path / "dump", "r+b")
-            self._mmap = mmap(self._file.fileno(), 0)
-            self._size = self.metadata.size
-            self._md5 = self.metadata.md5
-        return self
-
-
-class SharedObjectMetadata:
-    _id: str
-    _mempool: MemoryPool
-    _file: BinaryIO
-    _mmap: mmap
-    _instances: dict[str, SharedObjectMetadata] = {}
-
-    def __new__(cls, id: str, *, mempool: MemoryPool):
-        if id in cls._instances:
-            return cls._instances[id]
-        return super().__new__(cls)
-
-    def __init__(self, id: str, mempool: MemoryPool):
-        self._id = id
-        self._mempool = mempool
-        self._file = open(self._path / "meta", "r+b")
-        self._mmap = mmap(self._file.fileno(), 0)
-        self._instances[id] = self
-
-    def __del__(self):
-        try:
-            self.close()
-        except Exception:
-            pass
-
-    @property
-    def id(self) -> str:
-        return self._id
-
-    @property
-    def mutable(self) -> bool:
-        return self._metadata.mutable
-
-    @property
-    def size(self) -> int:
-        return self._metadata.size
-
-    @property
-    def md5(self) -> bytes:
-        return self._metadata.md5
-
-    @property
-    def mmap(self) -> mmap:
-        return self._mmap
-
-    @property
-    def _path(self) -> pathlib.Path:
-        return pathlib.Path(self._mempool.path, self.id)
-
-    @property
-    def _metadata(self) -> MetadataMessage:
-        return MetadataMessage.loads(bytes(self._mmap))
-
-    def close(self):
-        self._mmap.close()
-        self._file.close()
-        del self._instances[self.id]
-
-
-class MemoryPool:
-    _objects: dict[str, SharedObject]
-    _path: pathlib.Path
-
-    def __init__(self, path: str | pathlib.Path = pathlib.Path(".mempool")):
-        if isinstance(path, str):
-            self._path = pathlib.Path(path)
-        else:
-            self._path = path
-        self._lockdir = self._path / "locks"
-        os.makedirs(self._lockdir, exist_ok=True)
-        self._acquire(f"pid-{os.getpid()}")
-        self._objects = dict()
-
-    def __contains__(self, ref: str) -> bool:
-        return ref in self._objects
-
-    def __del__(self):
-        self._release(f"pid-{os.getpid()}")
-
-    @property
-    def path(self) -> pathlib.Path:
-        return self._path
-
-    async def map(self, ref: str | None = None):
-        if ref is not None and ref not in self._objects:
-            if self._locked(f"delete-{ref}"):
-                raise RuntimeError(
-                    f"Reference {ref} is currently locked for deletion"
-                )
-            async with self._reference_lock(ref):
-                self._map(ref)
-        else:
-            for entry in os.scandir(self._path):
-                if entry.is_dir() and (ref := entry.name) != "locks":
-                    if not self._locked(f"delete-{ref}"):
-                        async with self._reference_lock(ref):
-                            self._map(ref)
-
-    async def put(
-        self, dump: bytes, *, mutable: bool = False, ref: str | None = None
-    ) -> str:
-        ref = ref or str(shortuuid.uuid())
-        async with self._reference_lock(ref):
-            self._put(ref, dump, mutable=mutable, exist_ok=False)
-            return ref
-
-    async def post(self, ref: str, dump: bytes) -> bool:
-        if self._locked(f"delete-{ref}"):
-            raise RuntimeError(
-                f"Reference {ref} is currently locked for deletion"
-            )
-        async with self._reference_lock(ref):
-            if ref not in self._objects:
-                self._map(ref)
-            obj = self._objects[ref]
-            if not obj.metadata.mutable:
-                raise ValueError("Cannot modify an immutable reference")
-            if (size := len(dump)) != obj.metadata.size:
-                try:
-                    obj.mmap.resize(size)
-                    self._post(ref, obj, dump)
-                except SystemError:
-                    self._put(ref, dump, mutable=True, exist_ok=True)
-                return True
-            elif hashlib.md5(dump).digest() != obj.metadata.md5:
-                self._post(ref, obj, dump)
-                return True
-            else:
-                return False
-
-    async def get(self, ref: str) -> bytes:
-        if self._locked(f"delete-{ref}"):
-            raise RuntimeError(
-                f"Reference {ref} is currently locked for deletion"
-            )
-        async with self._reference_lock(ref):
-            if ref not in self._objects:
-                self._map(ref)
-            return bytes(self._objects[ref].refresh().mmap)
-
-    async def delete(self, ref: str):
-        async with self._delete_lock(ref):
-            async with self._reference_lock(ref):
-                if ref not in self._objects:
-                    self._map(ref)
-                self._objects.pop(ref).close()
-                try:
-                    shutil.rmtree(self.path / ref)
-                except FileNotFoundError:
-                    pass
-
-    def _put(
-        self,
-        ref: str,
-        dump: bytes,
-        *,
-        mutable: bool = False,
-        exist_ok: bool = False,
-    ):
-        metadata = MetadataMessage(
-            ref=ref,
-            mutable=mutable,
-            size=len(dump),
-            md5=hashlib.md5(dump).digest(),
-        )
-
-        refpath = pathlib.Path(self._path, f"{ref}")
-        os.makedirs(refpath, exist_ok=exist_ok)
-
-        with open(refpath / "meta", "wb") as metafile:
-            metafile.write(metadata.dumps())
-
-        with open(refpath / "dump", "wb") as dumpfile:
-            dumpfile.write(dump)
-
-        self._map(ref)
-
-    def _post(self, ref: str, obj: SharedObject, dump: bytes):
-        if not obj.metadata.mutable:
-            raise ValueError("Cannot modify an immutable reference")
-        metadata = MetadataMessage(
-            ref=ref,
-            mutable=True,
-            size=len(dump),
-            md5=hashlib.md5(dump).digest(),
-        )
-        obj.metadata.mmap[:] = metadata.dumps()
-        obj.metadata.mmap.flush()
-        obj.mmap.seek(0)
-        obj.mmap.write(dump)
-
-    def _map(self, ref: str):
-        obj = self._objects.pop(ref, None)
-        if obj:
-            obj.close()
-        self._objects[ref] = SharedObject(id=ref, mempool=self)
-
-    def _lockpath(self, key: str) -> pathlib.Path:
-        return pathlib.Path(self._lockdir, f"{key}.lock")
-
-    def _acquire(self, key: str) -> bool:
-        try:
-            os.symlink(f"{key}", self._lockpath(key))
-            return True
-        except FileExistsError:
-            return False
-
-    def _release(self, key: str):
-        try:
-            if os.path.islink(lock_path := self._lockpath(key)):
-                os.unlink(lock_path)
-        except FileNotFoundError:
-            pass
-
-    def _locked(self, key: str) -> bool:
-        return os.path.islink(self._lockpath(key))
-
-    @asynccontextmanager
-    async def _reference_lock(self, ref: str):
-        try:
-            while not self._acquire(ref):
-                await asyncio.sleep(0)
-            yield
-        finally:
-            self._release(ref)
-
-    @asynccontextmanager
-    async def _delete_lock(self, ref: str):
-        key = f"delete-{ref}"
-        if not self._acquire(f"delete-{ref}"):
-            raise RuntimeError(
-                f"Reference {ref} is currently locked for deletion"
-            )
-        else:
-            yield
-        self._release(key)

wool/_mempool/_metadata.py

@@ -1,35 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import asdict
-from dataclasses import dataclass
-
-try:
-    from wool._protobuf.mempool import metadata_pb2 as pb
-except ImportError as e:
-    from wool._protobuf import ProtobufImportError
-
-    raise ProtobufImportError(e) from e
-
-
-@dataclass
-class MetadataMessage:
-    ref: str
-    mutable: bool
-    size: int
-    md5: bytes
-
-    @classmethod
-    def loads(cls, data: bytes) -> MetadataMessage:
-        (metadata := pb.MetadataMessage()).ParseFromString(data)
-        return cls(
-            ref=metadata.ref,
-            mutable=metadata.mutable,
-            size=metadata.size,
-            md5=metadata.md5,
-        )
-
-    def dumps(self) -> bytes:
-        return pb.MetadataMessage(**asdict(self)).SerializeToString()
-
-
-__all__ = ["MetadataMessage"]