krons 0.1.0__py3-none-any.whl

kronos/utils/_json_dump.py

@@ -0,0 +1,392 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""JSON serialization utilities built on orjson.
+
+Provides flexible serialization with:
+- Configurable type handling (Decimal, Enum, datetime, sets)
+- Safe fallback mode for logging non-serializable objects
+- NDJSON streaming for iterables
+- Caching of default handlers for performance
+"""
+
+from __future__ import annotations
+
+import contextlib
+import datetime as dt
+import decimal
+import re
+from collections.abc import Callable, Iterable, Mapping
+from enum import Enum
+from functools import lru_cache
+from pathlib import Path
+from textwrap import shorten
+from typing import Any
+from uuid import UUID
+
+import orjson
+
+__all__ = (
+    "get_orjson_default",
+    "json_dumpb",
+    "json_dump",
+    "json_lines_iter",
+    "make_options",
+)
+
+# Types with native orjson support (skip in custom order)
+_NATIVE = (dt.datetime, dt.date, dt.time, UUID)
+_ADDR_PAT = re.compile(r" at 0x[0-9A-Fa-f]+")
+_SERIALIZATION_METHODS = ("model_dump", "to_dict", "dict")
+
+
+def get_orjson_default(
+    *,
+    order: list[type] | None = None,
+    additional: Mapping[type, Callable[[Any], Any]] | None = None,
+    extend_default: bool = True,
+    deterministic_sets: bool = False,
+    decimal_as_float: bool = False,
+    enum_as_name: bool = False,
+    passthrough_datetime: bool = False,
+    safe_fallback: bool = False,
+    fallback_clip: int = 2048,
+) -> Callable[[Any], Any]:
+    """Build a `default=` callable for orjson.dumps with type-based dispatch.
+
+    Args:
+        order: Custom type priority order (checked before defaults).
+        additional: Extra type->serializer mappings.
+        extend_default: Merge order with defaults (True) or replace (False).
+        deterministic_sets: Sort sets for reproducible output (slower).
+        decimal_as_float: Decimal->float (lossy but compact).
+        enum_as_name: Enum->name instead of value.
+        passthrough_datetime: Use custom datetime serialization.
+        safe_fallback: Never raise; clip repr for unknown types (for logging).
+        fallback_clip: Max chars for safe_fallback repr.
+
+    Returns:
+        Callable suitable for orjson.dumps(default=...).
+    """
+    ser = _default_serializers(
+        deterministic_sets=deterministic_sets,
+        decimal_as_float=decimal_as_float,
+        enum_as_name=enum_as_name,
+        passthrough_datetime=passthrough_datetime,
+    )
+    if additional:
+        ser.update(additional)
+
+    base_order: list[type] = [Path, decimal.Decimal, set, frozenset]
+    if enum_as_name:
+        base_order.insert(0, Enum)
+    if passthrough_datetime:
+        base_order.insert(0, dt.datetime)
+
+    if order:
+        order_ = (
+            (base_order + [t for t in order if t not in base_order])
+            if extend_default
+            else list(order)
+        )
+    else:
+        order_ = base_order.copy()
+
+    if not passthrough_datetime:  # Skip types on orjson's native fast path
+        order_ = [t for t in order_ if t not in _NATIVE]
+
+    order_tuple = tuple(order_)
+    cache: dict[type, Callable[[Any], Any]] = {}
+
+    def default(obj: Any) -> Any:
+        typ = obj.__class__
+        func = cache.get(typ)
+        if func is None:
+            for T in order_tuple:
+                if issubclass(typ, T):
+                    f = ser.get(T)
+                    if f:
+                        cache[typ] = f
+                        func = f
+                        break
+            else:
+                methods = _SERIALIZATION_METHODS
+                for m in methods:
+                    md = getattr(obj, m, None)
+                    if callable(md):
+                        with contextlib.suppress(Exception):
+                            return md()
+
+                if safe_fallback:
+                    if isinstance(obj, Exception):
+                        return {"type": obj.__class__.__name__, "message": str(obj)}
+
+                    return shorten(
+                        repr(obj),
+                        width=fallback_clip,
+                        placeholder=f"...(+{len(repr(obj)) - fallback_clip} chars)",
+                    )
+                raise TypeError(f"Type is not JSON serializable: {typ.__name__}")
+        return func(obj)
+
+    return default
+
+
+def make_options(
+    *,
+    pretty: bool = False,
+    sort_keys: bool = False,
+    naive_utc: bool = False,
+    utc_z: bool = False,
+    append_newline: bool = False,
+    passthrough_datetime: bool = False,
+    allow_non_str_keys: bool = False,
+) -> int:
+    """Compose orjson option bit flags.
+
+    Args:
+        pretty: Indent with 2 spaces (OPT_INDENT_2).
+        sort_keys: Alphabetical key ordering (OPT_SORT_KEYS).
+        naive_utc: Treat naive datetime as UTC (OPT_NAIVE_UTC).
+        utc_z: Use 'Z' suffix for UTC times (OPT_UTC_Z).
+        append_newline: Add trailing newline (OPT_APPEND_NEWLINE).
+        passthrough_datetime: Custom datetime handling (OPT_PASSTHROUGH_DATETIME).
+        allow_non_str_keys: Allow int/UUID keys (OPT_NON_STR_KEYS).
+
+    Returns:
+        Combined option flags for orjson.dumps(option=...).
+    """
+    opt = 0
+    if append_newline:
+        opt |= orjson.OPT_APPEND_NEWLINE
+    if pretty:
+        opt |= orjson.OPT_INDENT_2
+    if sort_keys:
+        opt |= orjson.OPT_SORT_KEYS
+    if naive_utc:
+        opt |= orjson.OPT_NAIVE_UTC
+    if utc_z:
+        opt |= orjson.OPT_UTC_Z
+    if passthrough_datetime:
+        opt |= orjson.OPT_PASSTHROUGH_DATETIME
+    if allow_non_str_keys:
+        opt |= orjson.OPT_NON_STR_KEYS
+    return opt
+
+
+def json_dumpb(
+    obj: Any,
+    *,
+    pretty: bool = False,
+    sort_keys: bool = False,
+    naive_utc: bool = False,
+    utc_z: bool = False,
+    append_newline: bool = False,
+    allow_non_str_keys: bool = False,
+    deterministic_sets: bool = False,
+    decimal_as_float: bool = False,
+    enum_as_name: bool = False,
+    passthrough_datetime: bool = False,
+    safe_fallback: bool = False,
+    fallback_clip: int = 2048,
+    default: Callable[[Any], Any] | None = None,
+    options: int | None = None,
+) -> bytes:
+    """Serialize to bytes (fast path for hot code).
+
+    Args:
+        obj: Object to serialize.
+        pretty: Indent output.
+        sort_keys: Alphabetical key order.
+        naive_utc: Naive datetime as UTC.
+        utc_z: Use 'Z' for UTC.
+        append_newline: Trailing newline.
+        allow_non_str_keys: Allow non-string dict keys.
+        deterministic_sets: Sort sets.
+        decimal_as_float: Decimal as float.
+        enum_as_name: Enum as name.
+        passthrough_datetime: Custom datetime handling.
+        safe_fallback: Never raise (for logging only).
+        fallback_clip: Max repr chars in safe mode.
+        default: Custom default callable (overrides above).
+        options: Pre-composed option flags (overrides above).
+
+    Returns:
+        JSON as bytes.
+    """
+    if default is None:
+        default = _cached_default(
+            deterministic_sets=deterministic_sets,
+            decimal_as_float=decimal_as_float,
+            enum_as_name=enum_as_name,
+            passthrough_datetime=passthrough_datetime,
+            safe_fallback=safe_fallback,
+            fallback_clip=fallback_clip,
+        )
+    opt = (
+        options
+        if options is not None
+        else make_options(
+            pretty=pretty,
+            sort_keys=sort_keys,
+            naive_utc=naive_utc,
+            utc_z=utc_z,
+            append_newline=append_newline,
+            passthrough_datetime=passthrough_datetime,
+            allow_non_str_keys=allow_non_str_keys,
+        )
+    )
+    return orjson.dumps(obj, default=default, option=opt)
+
+
+def json_dump(
+    obj: Any,
+    *,
+    sort_keys: bool = False,
+    deterministic_sets: bool = False,
+    decode: bool = False,
+    as_loaded: bool = False,
+    **kwargs: Any,
+) -> str | bytes | Any:
+    """Serialize to JSON with flexible output format.
+
+    Args:
+        obj: Object to serialize.
+        sort_keys: Alphabetical key order.
+        deterministic_sets: Sort sets.
+        decode: Return str instead of bytes.
+        as_loaded: Parse output back to dict/list (requires decode=True).
+        **kwargs: Passed to json_dumpb.
+
+    Returns:
+        bytes (default), str (decode=True), or dict/list (as_loaded=True).
+
+    Raises:
+        ValueError: If as_loaded=True without decode=True.
+    """
+    if not decode and as_loaded:
+        raise ValueError("as_loaded=True requires decode=True")
+
+    bytes_ = json_dumpb(
+        obj,
+        sort_keys=sort_keys,
+        deterministic_sets=deterministic_sets,
+        **kwargs,
+    )
+
+    if not decode:
+        return bytes_
+    return orjson.loads(bytes_) if as_loaded else bytes_.decode("utf-8")
+
+
+def json_lines_iter(
+    it: Iterable[Any],
+    *,
+    deterministic_sets: bool = False,
+    decimal_as_float: bool = False,
+    enum_as_name: bool = False,
+    passthrough_datetime: bool = False,
+    safe_fallback: bool = False,
+    fallback_clip: int = 2048,
+    naive_utc: bool = False,
+    utc_z: bool = False,
+    allow_non_str_keys: bool = False,
+    default: Callable[[Any], Any] | None = None,
+    options: int | None = None,
+) -> Iterable[bytes]:
+    """Stream iterable as NDJSON (newline-delimited JSON bytes).
+
+    Each item serialized to one line with trailing newline. Suitable for
+    streaming large datasets or log output.
+
+    Args:
+        it: Iterable of objects to serialize.
+        deterministic_sets: Sort sets.
+        decimal_as_float: Decimal as float.
+        enum_as_name: Enum as name.
+        passthrough_datetime: Custom datetime handling.
+        safe_fallback: Never raise (for logging).
+        fallback_clip: Max repr chars in safe mode.
+        naive_utc: Naive datetime as UTC.
+        utc_z: Use 'Z' for UTC.
+        allow_non_str_keys: Allow non-string dict keys.
+        default: Custom default callable.
+        options: Pre-composed option flags (newline always added).
+
+    Yields:
+        JSON bytes for each item with trailing newline.
+    """
+    if default is None:
+        default = _cached_default(
+            deterministic_sets=deterministic_sets,
+            decimal_as_float=decimal_as_float,
+            enum_as_name=enum_as_name,
+            passthrough_datetime=passthrough_datetime,
+            safe_fallback=safe_fallback,
+            fallback_clip=fallback_clip,
+        )
+    if options is None:
+        opt = make_options(
+            pretty=False,
+            sort_keys=False,
+            naive_utc=naive_utc,
+            utc_z=utc_z,
+            append_newline=True,
+            passthrough_datetime=passthrough_datetime,
+            allow_non_str_keys=allow_non_str_keys,
+        )
+    else:
+        opt = options | orjson.OPT_APPEND_NEWLINE  # Always enforce newline
+
+    for item in it:
+        yield orjson.dumps(item, default=default, option=opt)
+
+
+@lru_cache(maxsize=128)
+def _cached_default(
+    deterministic_sets: bool,
+    decimal_as_float: bool,
+    enum_as_name: bool,
+    passthrough_datetime: bool,
+    safe_fallback: bool,
+    fallback_clip: int,
+) -> Callable[[Any], Any]:
+    """Cache default handlers to avoid repeated construction."""
+    return get_orjson_default(
+        deterministic_sets=deterministic_sets,
+        decimal_as_float=decimal_as_float,
+        enum_as_name=enum_as_name,
+        passthrough_datetime=passthrough_datetime,
+        safe_fallback=safe_fallback,
+        fallback_clip=fallback_clip,
+    )
+
+
+def _default_serializers(
+    deterministic_sets: bool,
+    decimal_as_float: bool,
+    enum_as_name: bool,
+    passthrough_datetime: bool,
+) -> dict[type, Callable[[Any], Any]]:
+    """Build type->serializer mapping based on configuration."""
+
+    def normalize_for_sorting(x: Any) -> str:
+        """Strip memory addresses from repr for stable sorting."""
+        return _ADDR_PAT.sub(" at 0x?", str(x))
+
+    def stable_sorted_iterable(o: Iterable[Any]) -> list[Any]:
+        """Sort mixed-type iterables by (class_name, normalized_repr)."""
+        return sorted(o, key=lambda x: (x.__class__.__name__, normalize_for_sorting(x)))
+
+    ser: dict[type, Callable[[Any], Any]] = {
+        Path: str,
+        decimal.Decimal: (float if decimal_as_float else str),
+        set: (stable_sorted_iterable if deterministic_sets else list),
+        frozenset: (stable_sorted_iterable if deterministic_sets else list),
+    }
+    if enum_as_name:
+        ser[Enum] = lambda e: e.name
+
+    if passthrough_datetime:
+        ser[dt.datetime] = lambda o: o.isoformat()
+    return ser

kronos/utils/_lazy_init.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Thread-safe lazy initialization utility."""
+
+import threading
+from collections.abc import Callable
+
+__all__ = ("LazyInit",)
+
+
+class LazyInit:
+    """Thread-safe lazy initialization helper using double-checked locking.
+
+    Defers expensive imports/setup until first use. Guarantees init_func
+    runs exactly once even under concurrent access.
+
+    Example:
+        _lazy = LazyInit()
+        _MODEL_LIKE = None
+
+        def _do_init():
+            global _MODEL_LIKE
+            from pydantic import BaseModel
+            _MODEL_LIKE = (BaseModel,)
+
+        def my_function(x):
+            _lazy.ensure(_do_init)
+            # _MODEL_LIKE is now initialized
+
+    Attributes:
+        initialized: True after ensure() has completed successfully.
+    """
+
+    __slots__ = ("_initialized", "_lock")
+
+    def __init__(self) -> None:
+        """Create uninitialized LazyInit instance."""
+        self._initialized = False
+        self._lock = threading.RLock()
+
+    @property
+    def initialized(self) -> bool:
+        """Check if initialization has completed."""
+        return self._initialized
+
+    def ensure(self, init_func: Callable[[], None]) -> None:
+        """Execute init_func exactly once, thread-safely.
+
+        Uses double-checked locking: fast path (no lock) when already
+        initialized, lock acquisition only on first call.
+
+        Args:
+            init_func: Initialization function. Should be idempotent
+                as a safety measure (though only called once).
+        """
+        if self._initialized:
+            return
+        with self._lock:
+            if self._initialized:
+                return
+            init_func()
+            self._initialized = True

kronos/utils/_to_list.py

@@ -0,0 +1,165 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""List conversion utilities with flattening, deduplication, and NA handling."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from enum import Enum as _Enum
+from typing import Any
+
+from ._hash import compute_hash
+from ._lazy_init import LazyInit
+
+__all__ = ("to_list",)
+
+_lazy = LazyInit()
+_MODEL_LIKE = None
+_MAP_LIKE = None
+_SINGLETONE_TYPES = None
+_SKIP_TYPE = None
+_SKIP_TUPLE_SET = None
+_BYTE_LIKE = (str, bytes, bytearray)
+_TUPLE_SET = (tuple, set, frozenset)
+
+
+def _do_init() -> None:
+    """Initialize lazy type constants (Pydantic, kron sentinels)."""
+    from pydantic import BaseModel
+    from pydantic_core import PydanticUndefinedType
+
+    from kronos.types import UndefinedType, UnsetType
+
+    global _MODEL_LIKE, _MAP_LIKE, _SINGLETONE_TYPES, _SKIP_TYPE, _SKIP_TUPLE_SET
+    _MODEL_LIKE = (BaseModel,)
+    _MAP_LIKE = (Mapping, *_MODEL_LIKE)
+    _SINGLETONE_TYPES = (UndefinedType, UnsetType, PydanticUndefinedType)
+    _SKIP_TYPE = (*_BYTE_LIKE, *_MAP_LIKE, _Enum)
+    _SKIP_TUPLE_SET = (*_SKIP_TYPE, *_TUPLE_SET)
+
+
+def to_list(
+    input_: Any,
+    /,
+    *,
+    flatten: bool = False,
+    dropna: bool = False,
+    unique: bool = False,
+    use_values: bool = False,
+    flatten_tuple_set: bool = False,
+) -> list:
+    """Convert input to list with optional transformations.
+
+    Type handling:
+        - None/Undefined/Unset: returns []
+        - list: returned as-is (not copied)
+        - Enum class: list of members (or values if use_values=True)
+        - str/bytes/bytearray: wrapped as [input_] unless use_values=True
+        - Mapping: wrapped as [input_] unless use_values=True (extracts values)
+        - BaseModel: wrapped as [input_]
+        - Other iterables: converted via list()
+        - Non-iterables: wrapped as [input_]
+
+    Args:
+        input_: Value to convert.
+        flatten: Recursively flatten nested iterables.
+        dropna: Remove None and sentinel values (Undefined, Unset).
+        unique: Remove duplicates. Requires flatten=True.
+        use_values: Extract values from Enum classes and Mappings.
+        flatten_tuple_set: When flatten=True, also flatten tuples/sets/frozensets.
+
+    Returns:
+        Processed list.
+
+    Raises:
+        ValueError: unique=True without flatten=True, or unhashable non-mapping item.
+
+    Edge Cases:
+        - Nested lists: preserved unless flatten=True
+        - Unhashable items with unique=True: falls back to compute_hash for mappings
+        - Empty input: returns []
+    """
+    _lazy.ensure(_do_init)
+
+    def _process_list(
+        lst: list[Any],
+        flatten: bool,
+        dropna: bool,
+        skip_types: tuple[type, ...],
+    ) -> list[Any]:
+        """Recursively process list with flatten/dropna logic."""
+        result: list[Any] = []
+        for item in lst:
+            if dropna and (item is None or isinstance(item, _SINGLETONE_TYPES)):
+                continue
+            is_iterable = isinstance(item, Iterable)
+            should_skip = isinstance(item, skip_types)
+            if is_iterable and not should_skip:
+                item_list = list(item)
+                if flatten:
+                    result.extend(_process_list(item_list, flatten, dropna, skip_types))
+                else:
+                    result.append(_process_list(item_list, flatten, dropna, skip_types))
+            else:
+                result.append(item)
+        return result
+
+    def _to_list_type(input_: Any, use_values: bool) -> list[Any]:
+        """Convert input to initial list based on type."""
+        if input_ is None or isinstance(input_, _SINGLETONE_TYPES):
+            return []
+        if isinstance(input_, list):
+            return input_
+        if isinstance(input_, type) and issubclass(input_, _Enum):
+            members = input_.__members__.values()
+            return [member.value for member in members] if use_values else list(members)
+        if isinstance(input_, _BYTE_LIKE):
+            return list(input_) if use_values else [input_]
+        if isinstance(input_, Mapping):
+            return list(input_.values()) if use_values and hasattr(input_, "values") else [input_]
+        if isinstance(input_, _MODEL_LIKE):
+            return [input_]
+        if isinstance(input_, Iterable) and not isinstance(input_, _BYTE_LIKE):
+            return list(input_)
+        return [input_]
+
+    if unique and not flatten:
+        raise ValueError("unique=True requires flatten=True")
+
+    initial_list = _to_list_type(input_, use_values=use_values)
+    skip_types: tuple[type, ...] = _SKIP_TYPE if flatten_tuple_set else _SKIP_TUPLE_SET
+    processed = _process_list(initial_list, flatten=flatten, dropna=dropna, skip_types=skip_types)
+
+    if unique:
+        seen = set()
+        out = []
+        use_hash_fallback = False
+        for i in processed:
+            try:
+                if not use_hash_fallback and i not in seen:
+                    seen.add(i)
+                    out.append(i)
+            except TypeError:
+                if not use_hash_fallback:
+                    # Restart with hash-based deduplication
+                    use_hash_fallback = True
+                    seen = set()
+                    out = []
+                    for j in processed:
+                        try:
+                            hash_value = hash(j)
+                        except TypeError:
+                            if _MAP_LIKE is not None and isinstance(j, _MAP_LIKE):
+                                hash_value = compute_hash(j)
+                            else:
+                                raise ValueError(
+                                    "Unhashable type encountered in list unique value processing."
+                                )
+                        if hash_value not in seen:
+                            seen.add(hash_value)
+                            out.append(j)
+                    break
+        return out
+
+    return processed

kronos/utils/_to_num.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Numeric conversion utilities with bounds checking and security limits."""
+
+import math
+from decimal import Decimal
+from typing import Any
+
+__all__ = ("to_num", "MAX_NUMBER_STRING_LENGTH")
+
+MAX_NUMBER_STRING_LENGTH = 1000
+"""Max string length for numeric conversion (DoS protection)."""
+
+
+def to_num(
+    input_: Any,
+    /,
+    *,
+    upper_bound: int | float | None = None,
+    lower_bound: int | float | None = None,
+    num_type: type[int] | type[float] = float,
+    precision: int | None = None,
+    allow_inf: bool = False,
+) -> int | float:
+    """Convert input to numeric type with bounds checking and validation.
+
+    Handles: bool, int, float, Decimal, str. Strings are stripped before parsing.
+
+    Args:
+        input_: Value to convert. Bools treated as int (True=1, False=0).
+        upper_bound: Maximum allowed value (inclusive).
+        lower_bound: Minimum allowed value (inclusive).
+        num_type: Target type, must be `int` or `float`.
+        precision: Decimal places for float rounding (ignored for int).
+        allow_inf: Permit infinity values. Default False.
+
+    Returns:
+        Converted numeric value of type `num_type`.
+
+    Raises:
+        ValueError: Empty/too-long string, out of bounds, NaN, inf (when disallowed).
+        TypeError: Unsupported input type.
+
+    Edge Cases:
+        - Empty string: raises ValueError
+        - Whitespace-only string: raises ValueError (stripped to empty)
+        - "inf"/"-inf": raises unless allow_inf=True
+        - "nan": always raises
+        - Decimal: converted via float (may lose precision)
+    """
+    if num_type not in (int, float):
+        raise ValueError(f"Invalid number type: {num_type}")
+
+    if isinstance(input_, (bool, int, float, Decimal)):
+        value = float(input_)
+    elif isinstance(input_, str):
+        input_ = input_.strip()
+        if not input_:
+            raise ValueError("Empty string cannot be converted to number")
+        if len(input_) > MAX_NUMBER_STRING_LENGTH:
+            msg = f"String length ({len(input_)}) exceeds maximum ({MAX_NUMBER_STRING_LENGTH})"
+            raise ValueError(msg)
+        try:
+            value = float(input_)
+        except ValueError as e:
+            raise ValueError(f"Cannot convert '{input_}' to number") from e
+    else:
+        raise TypeError(f"Cannot convert {type(input_).__name__} to number")
+
+    # NaN bypasses all comparisons; always reject
+    if math.isnan(value):
+        raise ValueError("NaN is not allowed")
+    if math.isinf(value) and not allow_inf:
+        raise ValueError("Infinity is not allowed (use allow_inf=True to permit)")
+
+    if upper_bound is not None and value > upper_bound:
+        raise ValueError(f"Value {value} exceeds upper bound {upper_bound}")
+    if lower_bound is not None and value < lower_bound:
+        raise ValueError(f"Value {value} below lower bound {lower_bound}")
+
+    if precision is not None and num_type is float:
+        value = round(value, precision)
+
+    return num_type(value)