krons 0.1.0__py3-none-any.whl

kronos/utils/_utils.py

@@ -0,0 +1,375 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+import contextlib
+import importlib.util
+from collections.abc import Awaitable, Callable
+from datetime import UTC, datetime
+from functools import wraps
+from pathlib import Path as StdPath
+from types import UnionType
+from typing import Any, ParamSpec, TypeVar, Union, get_args, get_origin
+from uuid import UUID, uuid4
+
+from anyio import Path as AsyncPath
+
+from kronos.protocols import Observable
+
+__all__ = (
+    "create_path",
+    "get_bins",
+    "import_module",
+    "is_import_installed",
+    "now_utc",
+    "to_uuid",
+    "coerce_created_at",
+    "load_type_from_string",
+    "async_synchronized",
+    "synchronized",
+)
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+def now_utc() -> datetime:
+    """Get current UTC datetime."""
+    return datetime.now(UTC)
+
+
+async def create_path(
+    directory: StdPath | AsyncPath | str,
+    filename: str,
+    extension: str | None = None,
+    timestamp: bool = False,
+    dir_exist_ok: bool = True,
+    file_exist_ok: bool = False,
+    time_prefix: bool = False,
+    timestamp_format: str | None = None,
+    random_hash_digits: int = 0,
+    timeout: float | None = None,
+) -> AsyncPath:
+    """Generate file path asynchronously with optional timeout.
+
+    Args:
+        directory: Base directory path
+        filename: Target filename (may contain subdirectory with /)
+        extension: File extension (if filename doesn't have one)
+        timestamp: Add timestamp to filename
+        dir_exist_ok: Allow existing directories
+        file_exist_ok: Allow existing files
+        time_prefix: Put timestamp before filename instead of after
+        timestamp_format: Custom strftime format for timestamp
+        random_hash_digits: Add random hash suffix (0 = disabled)
+        timeout: Maximum time in seconds for async I/O operations (None = no timeout)
+
+    Returns:
+        AsyncPath to the created/validated file path
+
+    Raises:
+        ValueError: If filename contains backslash
+        FileExistsError: If file exists and file_exist_ok is False
+        TimeoutError: If timeout is exceeded
+    """
+    from .concurrency import move_on_after
+
+    async def _impl() -> AsyncPath:
+        nonlocal directory, filename
+
+        if "/" in filename:
+            sub_dir, filename = filename.split("/")[:-1], filename.split("/")[-1]
+            directory = AsyncPath(directory) / "/".join(sub_dir)
+
+        if "\\" in filename:
+            raise ValueError("Filename cannot contain directory separators.")
+
+        directory = AsyncPath(directory)
+        if "." in filename:
+            name, ext = filename.rsplit(".", 1)
+        else:
+            name = filename
+            ext = extension or ""
+        ext = f".{ext.lstrip('.')}" if ext else ""
+
+        if timestamp:
+            ts_str = datetime.now().strftime(timestamp_format or "%Y%m%d%H%M%S")
+            name = f"{ts_str}_{name}" if time_prefix else f"{name}_{ts_str}"
+
+        if random_hash_digits > 0:
+            random_suffix = uuid4().hex[:random_hash_digits]
+            name = f"{name}-{random_suffix}"
+
+        full_path = directory / f"{name}{ext}"
+
+        await full_path.parent.mkdir(parents=True, exist_ok=dir_exist_ok)
+
+        if await full_path.exists() and not file_exist_ok:
+            raise FileExistsError(f"File {full_path} already exists.")
+
+        return full_path
+
+    if timeout is None:
+        return await _impl()
+
+    with move_on_after(timeout) as cancel_scope:
+        result = await _impl()
+    if cancel_scope.cancelled_caught:
+        raise TimeoutError(f"create_path timed out after {timeout}s")
+    return result
+
+
+def get_bins(input_: list[str], upper: int) -> list[list[int]]:
+    """Partition string list indices into bins by cumulative length.
+
+    Args:
+        input_: List of strings to partition.
+        upper: Maximum cumulative length per bin.
+
+    Returns:
+        List of index lists, each bin's total string length < upper.
+    """
+    current = 0
+    bins = []
+    current_bin = []
+    for idx, item in enumerate(input_):
+        if current + len(item) < upper:
+            current_bin.append(idx)
+            current += len(item)
+        else:
+            bins.append(current_bin)
+            current_bin = [idx]
+            current = len(item)
+    if current_bin:
+        bins.append(current_bin)
+    return bins
+
+
+def import_module(
+    package_name: str,
+    module_name: str | None = None,
+    import_name: str | list | None = None,
+) -> Any:
+    """Import module or attribute(s) by dotted path.
+
+    Args:
+        package_name: Base package name (e.g., 'kron').
+        module_name: Optional submodule (e.g., 'utils').
+        import_name: Attribute(s) to import from module.
+
+    Returns:
+        Module object, single attribute, or list of attributes.
+
+    Raises:
+        ImportError: If module or attribute not found.
+    """
+    try:
+        full_import_path = f"{package_name}.{module_name}" if module_name else package_name
+
+        if import_name:
+            import_name = [import_name] if not isinstance(import_name, list) else import_name
+            a = __import__(
+                full_import_path,
+                fromlist=import_name,
+            )
+            if len(import_name) == 1:
+                return getattr(a, import_name[0])
+            return [getattr(a, name) for name in import_name]
+        else:
+            return __import__(full_import_path)
+
+    except ImportError as e:
+        raise ImportError(f"Failed to import module {full_import_path}: {e}") from e
+
+
+def is_import_installed(package_name: str) -> bool:
+    """Check if package is installed."""
+    return importlib.util.find_spec(package_name) is not None
+
+
+_TYPE_CACHE: dict[str, type] = {}
+
+_DEFAULT_ALLOWED_PREFIXES: frozenset[str] = frozenset({"kronos."})
+_ALLOWED_MODULE_PREFIXES: set[str] = set(_DEFAULT_ALLOWED_PREFIXES)
+
+
+def register_type_prefix(prefix: str) -> None:
+    """Register module prefix for dynamic type loading allowlist.
+
+    Security: Only register prefixes for modules you control.
+
+    Args:
+        prefix: Module prefix to allow (e.g., "myapp.models.").
+                Must end with "." to prevent prefix attacks.
+
+    Raises:
+        ValueError: If prefix doesn't end with ".".
+    """
+    if not prefix.endswith("."):
+        raise ValueError(f"Prefix must end with '.': {prefix}")
+    _ALLOWED_MODULE_PREFIXES.add(prefix)
+
+
+def load_type_from_string(type_str: str) -> type:
+    """Load type from fully qualified path (e.g., 'kronos.core.Node').
+
+    Security: Only allowlisted module prefixes can be loaded.
+
+    Args:
+        type_str: Fully qualified type path.
+
+    Returns:
+        The loaded type class.
+
+    Raises:
+        ValueError: If path invalid, not allowlisted, or type not found.
+    """
+    if type_str in _TYPE_CACHE:
+        return _TYPE_CACHE[type_str]
+
+    if not isinstance(type_str, str):
+        raise ValueError(f"Expected string, got {type(type_str)}")
+
+    if "." not in type_str:
+        raise ValueError(f"Invalid type path (no module): {type_str}")
+
+    if not any(type_str.startswith(prefix) for prefix in _ALLOWED_MODULE_PREFIXES):
+        raise ValueError(
+            f"Module '{type_str}' not in allowed prefixes: {sorted(_ALLOWED_MODULE_PREFIXES)}"
+        )
+
+    try:
+        module_path, class_name = type_str.rsplit(".", 1)
+        import importlib
+
+        module = importlib.import_module(module_path)
+        if module is None:
+            raise ImportError(f"Module '{module_path}' not found")
+
+        type_class = getattr(module, class_name)
+        if not isinstance(type_class, type):
+            raise ValueError(f"'{type_str}' is not a type")
+
+        _TYPE_CACHE[type_str] = type_class
+        return type_class
+
+    except (ValueError, ImportError, AttributeError) as e:
+        raise ValueError(f"Failed to load type '{type_str}': {e}") from e
+
+
+def extract_types(item_type: Any) -> set[type]:
+    """Extract concrete types from type annotations.
+
+    Handles Union, list, set, and single types recursively.
+
+    Args:
+        item_type: Type annotation (Union[X, Y], list[type], set[type], or type).
+
+    Returns:
+        Set of concrete types extracted from the annotation.
+    """
+
+    def is_union(t):
+        origin = get_origin(t)
+        return origin is Union or isinstance(t, UnionType)
+
+    extracted: set[type] = set()
+
+    if isinstance(item_type, set):
+        for t in item_type:
+            if is_union(t):
+                extracted.update(get_args(t))
+            else:
+                extracted.add(t)
+        return extracted
+
+    if isinstance(item_type, list):
+        for t in item_type:
+            if is_union(t):
+                extracted.update(get_args(t))
+            else:
+                extracted.add(t)
+        return extracted
+
+    if is_union(item_type):
+        return set(get_args(item_type))
+
+    return {item_type}
+
+
+def to_uuid(value: Any) -> UUID:
+    """Convert value to UUID instance.
+
+    Args:
+        value: UUID, UUID string, or Observable with .id attribute.
+
+    Returns:
+        UUID instance.
+
+    Raises:
+        ValueError: If value cannot be converted to UUID.
+    """
+    if isinstance(value, Observable):
+        return value.id
+    if isinstance(value, UUID):
+        return value
+    if isinstance(value, str):
+        return UUID(value)
+    raise ValueError("Cannot get ID from item.")
+
+
+def coerce_created_at(v) -> datetime:
+    """Coerce value to UTC-aware datetime.
+
+    Args:
+        v: datetime, Unix timestamp (int/float), or ISO string.
+
+    Returns:
+        UTC-aware datetime instance.
+
+    Raises:
+        ValueError: If value cannot be parsed as datetime.
+    """
+    if isinstance(v, datetime):
+        return v.replace(tzinfo=UTC) if v.tzinfo is None else v
+
+    if isinstance(v, (int, float)):
+        return datetime.fromtimestamp(v, tz=UTC)
+
+    if isinstance(v, str):
+        with contextlib.suppress(ValueError):
+            return datetime.fromtimestamp(float(v), tz=UTC)
+        with contextlib.suppress(ValueError):
+            return datetime.fromisoformat(v)
+        raise ValueError(f"String '{v}' is neither timestamp nor ISO format")
+
+    raise ValueError(f"Expected datetime/timestamp/string, got {type(v).__name__}")
+
+
+def synchronized(func: Callable[P, R]) -> Callable[P, R]:
+    """Decorator for thread-safe method execution.
+
+    Requires decorated method's instance to have self._lock (threading.Lock).
+    """
+
+    @wraps(func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        self = args[0]
+        with self._lock:
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+def async_synchronized(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
+    """Decorator for async-safe method execution.
+
+    Requires decorated method's instance to have self._async_lock (anyio.Lock).
+    """
+
+    @wraps(func)
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        self = args[0]
+        async with self._async_lock:  # type: ignore[attr-defined]
+            return await func(*args, **kwargs)
+
+    return wrapper

kronos/utils/concurrency/__init__.py

@@ -0,0 +1,205 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Async concurrency utilities with lazy loading for fast import.
+
+Core Patterns:
+    gather: Run awaitables concurrently, collect results in order.
+    race: Return first completion, cancel the rest.
+    bounded_map: Apply async function with concurrency limit.
+    retry: Exponential backoff with deadline awareness.
+    CompletionStream: Iterate results as they complete.
+
+Batch Processing:
+    alcall: Apply function to list with retry, timeout, concurrency control.
+    bcall: Batch wrapper yielding results per batch.
+
+Primitives (anyio wrappers):
+    Lock, Semaphore, Event, Condition, Queue, CapacityLimiter
+    TaskGroup, create_task_group
+
+Cancellation:
+    CancelScope, move_on_after, move_on_at, fail_after, fail_at
+    effective_deadline, get_cancelled_exc_class, is_cancelled
+
+Utilities:
+    run_async: Execute coroutine from sync context.
+    run_sync: Run sync function in thread pool.
+    sleep, current_time, is_coro_func
+
+Resource Tracking:
+    LeakTracker, track_resource, untrack_resource
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+_LAZY_IMPORTS: dict[str, tuple[str, str]] = {
+    # _cancel
+    "CancelScope": ("kronos.utils.concurrency._cancel", "CancelScope"),
+    "effective_deadline": ("kronos.utils.concurrency._cancel", "effective_deadline"),
+    "fail_after": ("kronos.utils.concurrency._cancel", "fail_after"),
+    "fail_at": ("kronos.utils.concurrency._cancel", "fail_at"),
+    "move_on_after": ("kronos.utils.concurrency._cancel", "move_on_after"),
+    "move_on_at": ("kronos.utils.concurrency._cancel", "move_on_at"),
+    # _errors
+    "get_cancelled_exc_class": (
+        "kronos.utils.concurrency._errors",
+        "get_cancelled_exc_class",
+    ),
+    "is_cancelled": ("kronos.utils.concurrency._errors", "is_cancelled"),
+    "non_cancel_subgroup": (
+        "kronos.utils.concurrency._errors",
+        "non_cancel_subgroup",
+    ),
+    "shield": ("kronos.utils.concurrency._errors", "shield"),
+    # _patterns
+    "CompletionStream": ("kronos.utils.concurrency._patterns", "CompletionStream"),
+    "bounded_map": ("kronos.utils.concurrency._patterns", "bounded_map"),
+    "gather": ("kronos.utils.concurrency._patterns", "gather"),
+    "race": ("kronos.utils.concurrency._patterns", "race"),
+    "retry": ("kronos.utils.concurrency._patterns", "retry"),
+    # _primitives
+    "CapacityLimiter": ("kronos.utils.concurrency._primitives", "CapacityLimiter"),
+    "Condition": ("kronos.utils.concurrency._primitives", "Condition"),
+    "Event": ("kronos.utils.concurrency._primitives", "Event"),
+    "Lock": ("kronos.utils.concurrency._primitives", "Lock"),
+    "Queue": ("kronos.utils.concurrency._primitives", "Queue"),
+    "Semaphore": ("kronos.utils.concurrency._primitives", "Semaphore"),
+    # _priority_queue
+    "PriorityQueue": ("kronos.utils.concurrency._priority_queue", "PriorityQueue"),
+    "QueueEmpty": ("kronos.utils.concurrency._priority_queue", "QueueEmpty"),
+    "QueueFull": ("kronos.utils.concurrency._priority_queue", "QueueFull"),
+    # _resource_tracker
+    "LeakInfo": ("kronos.utils.concurrency._resource_tracker", "LeakInfo"),
+    "LeakTracker": ("kronos.utils.concurrency._resource_tracker", "LeakTracker"),
+    "track_resource": (
+        "kronos.utils.concurrency._resource_tracker",
+        "track_resource",
+    ),
+    "untrack_resource": (
+        "kronos.utils.concurrency._resource_tracker",
+        "untrack_resource",
+    ),
+    # _run_async
+    "run_async": ("kronos.utils.concurrency._run_async", "run_async"),
+    # _task
+    "TaskGroup": ("kronos.utils.concurrency._task", "TaskGroup"),
+    "create_task_group": ("kronos.utils.concurrency._task", "create_task_group"),
+    # _utils
+    "current_time": ("kronos.utils.concurrency._utils", "current_time"),
+    "is_coro_func": ("kronos.utils.concurrency._utils", "is_coro_func"),
+    "run_sync": ("kronos.utils.concurrency._utils", "run_sync"),
+    "sleep": ("kronos.utils.concurrency._utils", "sleep"),
+    "alcall": ("kronos.utils.concurrency._async_call", "alcall"),
+    "bcall": ("kronos.utils.concurrency._async_call", "bcall"),
+}
+
+_LOADED: dict[str, object] = {}
+
+# Re-export built-in ExceptionGroup
+ExceptionGroup = ExceptionGroup
+
+
+def __getattr__(name: str) -> object:
+    """Lazy import attributes on first access."""
+    if name in _LOADED:
+        return _LOADED[name]
+
+    if name in _LAZY_IMPORTS:
+        from importlib import import_module
+
+        module_name, attr_name = _LAZY_IMPORTS[name]
+        module = import_module(module_name)
+        value = getattr(module, attr_name)
+        _LOADED[name] = value
+        return value
+
+    # Special case: ConcurrencyEvent is alias for Event
+    if name == "ConcurrencyEvent":
+        value = __getattr__("Event")
+        _LOADED[name] = value
+        return value
+
+    raise AttributeError(f"module 'kronos.utils.concurrency' has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    """Return all available attributes for autocomplete."""
+    return list(__all__)
+
+
+# TYPE_CHECKING block for static analysis
+if TYPE_CHECKING:
+    from ._async_call import alcall, bcall
+    from ._cancel import (
+        CancelScope,
+        effective_deadline,
+        fail_after,
+        fail_at,
+        move_on_after,
+        move_on_at,
+    )
+    from ._errors import (
+        get_cancelled_exc_class,
+        is_cancelled,
+        non_cancel_subgroup,
+        shield,
+    )
+    from ._patterns import CompletionStream, bounded_map, gather, race, retry
+    from ._primitives import CapacityLimiter, Condition, Event, Lock, Queue, Semaphore
+    from ._priority_queue import PriorityQueue, QueueEmpty, QueueFull
+    from ._resource_tracker import (
+        LeakInfo,
+        LeakTracker,
+        track_resource,
+        untrack_resource,
+    )
+    from ._run_async import run_async
+    from ._task import TaskGroup, create_task_group
+    from ._utils import current_time, is_coro_func, run_sync, sleep
+
+    ConcurrencyEvent = Event
+
+__all__ = (
+    "CancelScope",
+    "CapacityLimiter",
+    "CompletionStream",
+    "ConcurrencyEvent",
+    "Condition",
+    "Event",
+    "ExceptionGroup",
+    "LeakInfo",
+    "LeakTracker",
+    "Lock",
+    "PriorityQueue",
+    "Queue",
+    "QueueEmpty",
+    "QueueFull",
+    "Semaphore",
+    "TaskGroup",
+    "alcall",
+    "bcall",
+    "bounded_map",
+    "create_task_group",
+    "current_time",
+    "effective_deadline",
+    "fail_after",
+    "fail_at",
+    "gather",
+    "get_cancelled_exc_class",
+    "is_cancelled",
+    "is_coro_func",
+    "move_on_after",
+    "move_on_at",
+    "non_cancel_subgroup",
+    "race",
+    "retry",
+    "run_async",
+    "run_sync",
+    "shield",
+    "sleep",
+    "track_resource",
+    "untrack_resource",
+)