tmpkit 1.0.0__py3-none-any.whl

tmpkit/_decorators.py

@@ -0,0 +1,271 @@
+"""Decorators: @temp_dir() and @temp_file() for functions and test classes."""
+
+from __future__ import annotations
+
+import functools
+import inspect
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from tmpkit._async import temp_dir as async_temp_dir
+from tmpkit._async import temp_file as async_temp_file
+from tmpkit._sync import temp_dir as sync_temp_dir
+from tmpkit._sync import temp_file as sync_temp_file
+from tmpkit._types import StrPath
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def temp_dir(
+    *,
+    suffix: str | None = None,
+    prefix: str | None = None,
+    dir: StrPath | None = None,
+    cwd: bool = True,
+    keep: bool = False,
+    keep_on_error: bool = False,
+    ignore_cleanup_errors: bool = True,
+) -> Callable[[F], F]:
+    """Decorator that provides a fresh temp directory for each call.
+
+    By default ``cwd=True`` — the function runs inside the temp dir and
+    the original cwd is restored on exit.
+
+    Works on:
+    - Sync functions: wrapped normally.
+    - Async functions: detected via ``inspect.iscoroutinefunction``.
+    - Test classes: each method gets a fresh temp dir as ``self.tmpdir``.
+
+    Args:
+        suffix: Directory name suffix.
+        prefix: Directory name prefix.
+        dir: Parent directory. Defaults to system temp dir.
+        cwd: If ``True``, changes working directory to temp dir. Defaults to ``True``.
+        keep: If ``True``, directory is NOT removed.
+        keep_on_error: If ``True``, directory is kept only on exception.
+        ignore_cleanup_errors: If ``True``, ``OSError`` during cleanup is silently ignored.
+    """
+
+    def decorator(func_or_cls: F) -> F:
+        if inspect.isclass(func_or_cls):
+            _decorate_class(
+                func_or_cls,
+                suffix=suffix,
+                prefix=prefix,
+                dir=dir,
+                cwd=cwd,
+                keep=keep,
+                keep_on_error=keep_on_error,
+                ignore_cleanup_errors=ignore_cleanup_errors,
+            )
+            return func_or_cls
+
+        if inspect.iscoroutinefunction(func_or_cls):
+
+            @functools.wraps(func_or_cls)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                async with async_temp_dir(
+                    suffix=suffix,
+                    prefix=prefix,
+                    dir=dir,
+                    cwd=cwd,
+                    keep=keep,
+                    keep_on_error=keep_on_error,
+                    ignore_cleanup_errors=ignore_cleanup_errors,
+                ) as tmp:
+                    return await func_or_cls(tmp, *args, **kwargs)
+
+            return async_wrapper  # type: ignore[return-value]
+
+        @functools.wraps(func_or_cls)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with sync_temp_dir(
+                suffix=suffix,
+                prefix=prefix,
+                dir=dir,
+                cwd=cwd,
+                keep=keep,
+                keep_on_error=keep_on_error,
+                ignore_cleanup_errors=ignore_cleanup_errors,
+            ) as tmp:
+                return func_or_cls(tmp, *args, **kwargs)
+
+        return sync_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def temp_file(
+    *,
+    suffix: str | None = None,
+    prefix: str | None = None,
+    dir: StrPath | None = None,
+    mode: str = "w+b",
+    content: str | bytes | None = None,
+    dest: StrPath | None = None,
+    keep: bool = False,
+    keep_on_error: bool = False,
+    ignore_cleanup_errors: bool = True,
+) -> Callable[[F], F]:
+    """Decorator that injects a temp file as the first positional argument.
+
+    Works on sync and async **plain functions**. The temp file object is
+    injected as the very first positional argument, before any user-supplied
+    arguments. This decorator does **not** support methods or classes — use
+    it on module-level functions whose signature begins with the temp file
+    parameter.
+
+    Args:
+        suffix: File name suffix.
+        prefix: File name prefix.
+        dir: Parent directory. Defaults to system temp dir.
+        mode: Open mode. Defaults to ``"w+b"``.
+        content: Pre-populate file with this content.
+        dest: Destination path. On success, temp is moved here.
+        keep: If ``True``, file is NOT deleted.
+        keep_on_error: If ``True``, file is kept only on exception.
+        ignore_cleanup_errors: If ``True``, ``OSError`` during cleanup is silently ignored.
+    """
+
+    def decorator(func: F) -> F:
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                async with async_temp_file(
+                    suffix=suffix,
+                    prefix=prefix,
+                    dir=dir,
+                    mode=mode,
+                    content=content,
+                    dest=dest,
+                    keep=keep,
+                    keep_on_error=keep_on_error,
+                    ignore_cleanup_errors=ignore_cleanup_errors,
+                ) as f:
+                    return await func(f, *args, **kwargs)
+
+            return async_wrapper  # type: ignore[return-value]
+
+        @functools.wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with sync_temp_file(
+                suffix=suffix,
+                prefix=prefix,
+                dir=dir,
+                mode=mode,
+                content=content,
+                dest=dest,
+                keep=keep,
+                keep_on_error=keep_on_error,
+                ignore_cleanup_errors=ignore_cleanup_errors,
+            ) as f:
+                return func(f, *args, **kwargs)
+
+        return sync_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def _decorate_class(
+    cls: type,
+    *,
+    suffix: str | None,
+    prefix: str | None,
+    dir: StrPath | None,
+    cwd: bool,
+    keep: bool,
+    keep_on_error: bool,
+    ignore_cleanup_errors: bool,
+) -> None:
+    """Decorate a test class so each test method gets a fresh temp dir.
+
+    The temp dir path is available as ``self.tmpdir`` during the test.
+    Each method starting with ``test_`` is wrapped to enter/exit a temp dir
+    context around the original call.
+    """
+    for name, method in list(vars(cls).items()):
+        if not name.startswith("test_"):
+            continue
+        if not callable(method):
+            continue
+
+        if inspect.iscoroutinefunction(method):
+            wrapped = _wrap_async_test_method(
+                method,
+                suffix=suffix,
+                prefix=prefix,
+                dir=dir,
+                cwd=cwd,
+                keep=keep,
+                keep_on_error=keep_on_error,
+                ignore_cleanup_errors=ignore_cleanup_errors,
+            )
+        else:
+            wrapped = _wrap_sync_test_method(
+                method,
+                suffix=suffix,
+                prefix=prefix,
+                dir=dir,
+                cwd=cwd,
+                keep=keep,
+                keep_on_error=keep_on_error,
+                ignore_cleanup_errors=ignore_cleanup_errors,
+            )
+        setattr(cls, name, wrapped)
+
+
+def _wrap_sync_test_method(
+    method: Callable[..., Any],
+    *,
+    suffix: str | None,
+    prefix: str | None,
+    dir: StrPath | None,
+    cwd: bool,
+    keep: bool,
+    keep_on_error: bool,
+    ignore_cleanup_errors: bool,
+) -> Callable[..., Any]:
+    @functools.wraps(method)
+    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        with sync_temp_dir(
+            suffix=suffix,
+            prefix=prefix,
+            dir=dir,
+            cwd=cwd,
+            keep=keep,
+            keep_on_error=keep_on_error,
+            ignore_cleanup_errors=ignore_cleanup_errors,
+        ) as tmp:
+            self.tmpdir = tmp
+            return method(self, *args, **kwargs)
+
+    return wrapper
+
+
+def _wrap_async_test_method(
+    method: Callable[..., Any],
+    *,
+    suffix: str | None,
+    prefix: str | None,
+    dir: StrPath | None,
+    cwd: bool,
+    keep: bool,
+    keep_on_error: bool,
+    ignore_cleanup_errors: bool,
+) -> Callable[..., Any]:
+    @functools.wraps(method)
+    async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        async with async_temp_dir(
+            suffix=suffix,
+            prefix=prefix,
+            dir=dir,
+            cwd=cwd,
+            keep=keep,
+            keep_on_error=keep_on_error,
+            ignore_cleanup_errors=ignore_cleanup_errors,
+        ) as tmp:
+            self.tmpdir = tmp
+            return await method(self, *args, **kwargs)
+
+    return wrapper

tmpkit/_registry.py

@@ -0,0 +1,159 @@
+"""Temp registry: track all temp files/dirs for cleanup and inspection."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import threading
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+
+@dataclass(slots=True)
+class TempRecord:
+    """Record of a single temp file or directory.
+
+    Attributes:
+        path: Filesystem path of the temp.
+        kind: ``"file"`` or ``"dir"``.
+        created_at: Creation timestamp (epoch seconds from ``time.time``).
+        cleaned: ``True`` once the temp has been deleted/moved.
+        kept: ``True`` if the temp was kept (not cleaned up).
+    """
+
+    path: Path
+    kind: Literal["file", "dir"]
+    created_at: float
+    cleaned: bool = False
+    kept: bool = False
+
+
+class TempRegistry:
+    """Thread-safe registry of temporary files and directories.
+
+    Disabled by default for zero overhead. Enable via ``enable()`` or
+    setting ``TMPKIT_REGISTRY=1`` environment variable at import time.
+    """
+
+    __slots__ = ("_enabled", "_lock", "_records")
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._records: list[TempRecord] = []
+        self._enabled = os.environ.get("TMPKIT_REGISTRY") == "1"
+
+    @property
+    def enabled(self) -> bool:
+        """Whether the registry is currently active."""
+        with self._lock:
+            return self._enabled
+
+    def enable(self) -> None:
+        """Enable the registry."""
+        with self._lock:
+            self._enabled = True
+
+    def disable(self) -> None:
+        """Disable the registry. Existing records are preserved."""
+        with self._lock:
+            self._enabled = False
+
+    def register(self, path: Path, kind: Literal["file", "dir"]) -> TempRecord | None:
+        """Register a new temp. Returns the record, or ``None`` if disabled."""
+        with self._lock:
+            if not self._enabled:
+                return None
+            record = TempRecord(
+                path=path,
+                kind=kind,
+                created_at=_now(),
+            )
+            self._records.append(record)
+            return record
+
+    def mark_cleaned(self, record: TempRecord) -> None:
+        """Mark a record as cleaned up."""
+        with self._lock:
+            record.cleaned = True
+
+    def mark_kept(self, record: TempRecord) -> None:
+        """Mark a record as kept (not cleaned up)."""
+        with self._lock:
+            record.kept = True
+
+    @property
+    def all(self) -> list[TempRecord]:
+        """All records (active + cleaned)."""
+        with self._lock:
+            return list(self._records)
+
+    @property
+    def active(self) -> list[TempRecord]:
+        """Records for temps that have not been cleaned or kept."""
+        with self._lock:
+            return [r for r in self._records if not r.cleaned and not r.kept]
+
+    @property
+    def cleaned(self) -> list[TempRecord]:
+        """Records for temps that have been cleaned up."""
+        with self._lock:
+            return [r for r in self._records if r.cleaned]
+
+    def cleanup_all(self) -> int:
+        """Delete all active temps. Returns the number of temps deleted."""
+        count = 0
+        with self._lock:
+            active = [r for r in self._records if not r.cleaned and not r.kept]
+
+        for record in active:
+            # Re-check under the lock that the record is still active;
+            # another thread may have marked it kept/cleaned since the snapshot.
+            with self._lock:
+                if record.cleaned or record.kept:
+                    continue
+            try:
+                if not record.path.exists():
+                    # Already gone on disk — mark cleaned without counting.
+                    with self._lock:
+                        record.cleaned = True
+                    continue
+                if record.kind == "dir":
+                    shutil.rmtree(record.path)
+                else:
+                    record.path.unlink()
+                with self._lock:
+                    record.cleaned = True
+                count += 1
+            except OSError:
+                pass
+        return count
+
+    def keep_all(self) -> int:
+        """Mark all active temps as kept. Returns the number marked."""
+        with self._lock:
+            count = 0
+            for record in self._records:
+                if not record.cleaned and not record.kept:
+                    record.kept = True
+                    count += 1
+            return count
+
+    def clear_history(self) -> None:
+        """Remove all cleaned records from history."""
+        with self._lock:
+            self._records = [r for r in self._records if not r.cleaned]
+
+    def reset(self) -> None:
+        """Clear all records and disable. Useful for testing."""
+        with self._lock:
+            self._records.clear()
+            self._enabled = False
+
+
+def _now() -> float:
+    return time.time()
+
+
+temp_registry = TempRegistry()