diff --git a/packages/unxts.parametric/src/unxts/parametric/_src/config.py b/packages/unxts.parametric/src/unxts/parametric/_src/config.py index 49d14169..ee120084 100644 --- a/packages/unxts.parametric/src/unxts/parametric/_src/config.py +++ b/packages/unxts.parametric/src/unxts/parametric/_src/config.py @@ -16,20 +16,18 @@ "config", ) -import contextlib import tomllib from pathlib import Path from typing import Any, ClassVar, Final from traitlets import Bool, TraitError -from traitlets.config import Config, SingletonConfigurable +from traitlets.config import SingletonConfigurable from unxt._src.config import ( AbstractUnxtConfig, LocalConfigurable, _find_pyproject, _load_toml_config_from_pyproject, - _NestedConfigContext, ) PARAMETRIC_REPR_CONFIG_KEYS: Final = frozenset({"include_params"}) @@ -41,100 +39,15 @@ } -class _ParametricLocalConfig(LocalConfigurable): - """Base for the parametric repr/str configs with thread-local overrides. +class ParametricQuantityReprConfig(LocalConfigurable): + """``include_params`` for ``ParametricQuantity.__repr__`` (default ``False``). - Subclasses set ``_config_keys`` (their overridable trait names) and their - own ``_local`` thread-local storage; this base provides the generic - ``__getattribute__`` override lookup and ``override()`` context manager. + ``__getattribute__`` (thread-local override lookup) and ``override()`` (the + context manager) are inherited from + :class:`~unxt._src.config.LocalConfigurable`; this class only declares its + overridable trait and its ``_config_keys`` allowlist. """ - _config_keys: ClassVar[frozenset[str]] = frozenset() - - def __getattribute__(self, name: str) -> Any: - """Get attribute, checking thread-local overrides first.""" - keys = object.__getattribute__(self, "_config_keys") - if name in keys: - with contextlib.suppress(AttributeError): - local = object.__getattribute__(self, "_local") - if hasattr(local, "stack") and local.stack: - for overrides in reversed(local.stack): - if name in overrides: - return overrides[name] - - return object.__getattribute__(self, name) - - def override( - self, cfg: Config | None = None, /, **kwargs: Any - ) -> _NestedConfigContext: - """Create a context manager for temporary config changes. - - Examples - -------- - >>> import unxts.parametric as up - >>> with up.config.quantity_repr.override(include_params=True): - ... print(up.config.quantity_repr.include_params) - True - - """ - keys = self._config_keys - cls_name = type(self).__name__ - - if cfg is not None and kwargs: - msg = "Cannot specify both cfg and keyword arguments to override()" - raise ValueError(msg) - - if kwargs: - unknown_keys = set(kwargs) - keys - if unknown_keys: - valid_keys = ", ".join(sorted(keys)) - unknown = ", ".join(sorted(unknown_keys)) - msg = ( - f"Unknown {cls_name} override option(s): {unknown}. " - f"Valid options are: {valid_keys}" - ) - raise ValueError(msg) - - # Validate/resolve through traitlets immediately for fast, clear errors. - temp_instance = self.__class__() - validated_kwargs: dict[str, Any] = {} - for key, value in kwargs.items(): - try: - setattr(temp_instance, key, value) - except TraitError as e: - msg = ( - f"Invalid value for {cls_name} override option " - f"'{key}': {value!r}" - ) - raise ValueError(msg) from e - validated_kwargs[key] = object.__getattribute__(temp_instance, key) - kwargs = validated_kwargs - - if cfg is not None: - temp_instance = self.__class__(config=cfg) - # Only override the traits the Config actually specifies. Iterating - # every trait would read the class default for traits absent from - # ``cfg`` and clobber any globally-configured value on entry. - specified: set[str] = set() - for klass in type(self).__mro__: - section = cfg.get(klass.__name__, None) - if section: - specified |= set(section) - # Intersect with the *override allowlist*, not all traits: - # ``__getattribute__`` only consults these keys, so an entry outside - # it could never be observed -- it would just be a dead push. - overrides = { - name: object.__getattribute__(temp_instance, name) - for name in specified & set(self._config_keys) - } - kwargs = overrides - - return _NestedConfigContext(self, kwargs) - - -class ParametricQuantityReprConfig(_ParametricLocalConfig): - """``include_params`` for ``ParametricQuantity.__repr__`` (default ``False``).""" - _config_keys: ClassVar[frozenset[str]] = PARAMETRIC_REPR_CONFIG_KEYS include_params: ClassVar[object] = Bool( @@ -143,8 +56,12 @@ class ParametricQuantityReprConfig(_ParametricLocalConfig): ).tag(config=True) -class ParametricQuantityStrConfig(_ParametricLocalConfig): - """``include_params`` for ``ParametricQuantity.__str__`` (default ``True``).""" +class ParametricQuantityStrConfig(LocalConfigurable): + """``include_params`` for ``ParametricQuantity.__str__`` (default ``True``). + + See :class:`ParametricQuantityReprConfig` — the override machinery is + inherited from :class:`~unxt._src.config.LocalConfigurable`. + """ _config_keys: ClassVar[frozenset[str]] = PARAMETRIC_STR_CONFIG_KEYS diff --git a/src/unxt/_src/config.py b/src/unxt/_src/config.py index 6a5b5b9d..138e3f9e 100644 --- a/src/unxt/_src/config.py +++ b/src/unxt/_src/config.py @@ -31,11 +31,15 @@ class LocalConfigurable(Configurable): This class provides a mechanism for temporary, thread-local configuration changes that can be used in nested contexts (e.g., within a quantity's - ``__repr__()``). + ``__repr__()``). Concrete subclasses only declare their traits and set + ``_config_keys`` (the trait names eligible for override); the override + lookup and the ``override()`` context manager live here. """ # Thread-local storage for context manager overrides _local: threading.local + # Names of the config traits eligible for thread-local override + _config_keys: ClassVar[frozenset[str]] = frozenset() def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) @@ -45,114 +49,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # every other instance, including the global ``unxt.config`` singleton. object.__setattr__(self, "_local", threading.local()) - -@dataclass(frozen=True, slots=True) -class _NestedConfigContext: - """Context manager for temporary config changes on nested config objects. - - This class should not be instantiated directly. Use the override() method - on nested config objects (e.g., config.quantity_repr.override(...)). - """ - - config: LocalConfigurable # The nested config object (e.g., QuantityReprConfig) - overrides: dict[str, Any] # {attr: value} - - def __enter__(self) -> LocalConfigurable: - """Enter context, pushing overrides onto thread-local stack.""" - # Initialize thread-local stack if needed - if not hasattr(self.config._local, "stack"): # noqa: SLF001 - self.config._local.stack = [] # noqa: SLF001 - - # Push overrides onto the stack - self.config._local.stack.append(self.overrides) # noqa: SLF001 - return self.config - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb: TracebackType | None, - ) -> None: - """Exit context, popping overrides from thread-local stack.""" - # Pop overrides from the stack - if hasattr(self.config._local, "stack") and self.config._local.stack: # noqa: SLF001 - self.config._local.stack.pop() # noqa: SLF001 - - -# ============================================================================ -# Quantity `__repr__` - -QUANTITY_REPR_CONFIG_KEYS: Final = frozenset( - {"short_arrays", "use_short_name", "named_unit", "indent"} -) - - -class QuantityReprConfig(LocalConfigurable): - """Configuration for quantity ``__repr__()`` display options. - - This controls how quantity objects are displayed in ``repr()``. It is - consumed by ``AbstractQuantity.__repr__`` and so applies to every quantity - class, including the default ``Quantity`` (and ``ParametricQuantity`` from - the ``unxts.parametric`` package). - - Attributes - ---------- - short_arrays : bool | Literal["compact"] - Controls how arrays are displayed in repr. Options: - - "compact": Show array values without Array wrapper - - `True`: Show short array summary (shape/dtype) - - `False`: Show full array representation - Default: `False` - use_short_name : bool - If True and a class has a `short_name` attribute, use the short - name instead of the full class name in repr. Default: False - named_unit : bool - If `True`, display unit as a named argument `unit='m'`. - If `False`, display unit as a positional argument `'m'`. - Default: `True` - - Examples - -------- - >>> import unxt as u - >>> with u.config.quantity_repr.override( - ... short_arrays="compact", use_short_name=True - ... ): - ... q = u.Q([1, 2, 3], "m") - ... print(repr(q)) - Q([1, 2, 3], unit='m') - - """ - - short_arrays: ClassVar[object] = Union( - [Bool(), Enum(("compact",))], - default_value=False, - help=( - "Controls array display in repr. " - "Options: 'compact' (values only), " - "True (shape/dtype), False (full repr)" - ), - ).tag(config=True) - - use_short_name: ClassVar[object] = Bool( - default_value=False, - help=( - "Use a class's short name if available (e.g. Quantity -> Q, " - "ParametricQuantity -> PQ)" - ), - ).tag(config=True) - - named_unit: ClassVar[object] = Bool( - default_value=True, - help="Display unit as named argument (unit='m') vs positional ('m')", - ).tag(config=True) - - indent: ClassVar[object] = Int( - default_value=4, help="Indentation level for nested structures in repr" - ).tag(config=True) - def __getattribute__(self, name: str) -> Any: """Get attribute, checking thread-local overrides first.""" - if name in QUANTITY_REPR_CONFIG_KEYS: + if name in object.__getattribute__(self, "_config_keys"): with contextlib.suppress(AttributeError): local = object.__getattribute__(self, "_local") # Return the most recent override for this attribute @@ -216,17 +115,20 @@ def override( Quantity([1, 2, 3], unit='m') """ + keys = self._config_keys + cls_name = type(self).__name__ + if cfg is not None and kwargs: msg = "Cannot specify both cfg and keyword arguments to override()" raise ValueError(msg) if kwargs: - unknown_keys = set(kwargs) - QUANTITY_REPR_CONFIG_KEYS + unknown_keys = set(kwargs) - keys if unknown_keys: - valid_keys = ", ".join(sorted(QUANTITY_REPR_CONFIG_KEYS)) + valid_keys = ", ".join(sorted(keys)) unknown = ", ".join(sorted(unknown_keys)) msg = ( - f"Unknown QuantityReprConfig override option(s): {unknown}. " + f"Unknown {cls_name} override option(s): {unknown}. " f"Valid options are: {valid_keys}" ) raise ValueError(msg) @@ -240,7 +142,7 @@ def override( setattr(temp_instance, key, value) except TraitError as e: msg = ( - "Invalid value for QuantityReprConfig override option " + f"Invalid value for {cls_name} override option " f"'{key}': {value!r}" ) raise ValueError(msg) from e @@ -267,13 +169,120 @@ def override( # it could never be observed -- it would just be a dead push. overrides = { name: object.__getattribute__(temp_instance, name) - for name in specified & QUANTITY_REPR_CONFIG_KEYS + for name in specified & keys } kwargs = overrides return _NestedConfigContext(self, kwargs) +@dataclass(frozen=True, slots=True) +class _NestedConfigContext: + """Context manager for temporary config changes on nested config objects. + + This class should not be instantiated directly. Use the override() method + on nested config objects (e.g., config.quantity_repr.override(...)). + """ + + config: LocalConfigurable # The nested config object (e.g., QuantityReprConfig) + overrides: dict[str, Any] # {attr: value} + + def __enter__(self) -> LocalConfigurable: + """Enter context, pushing overrides onto thread-local stack.""" + # Initialize thread-local stack if needed + if not hasattr(self.config._local, "stack"): # noqa: SLF001 + self.config._local.stack = [] # noqa: SLF001 + + # Push overrides onto the stack + self.config._local.stack.append(self.overrides) # noqa: SLF001 + return self.config + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + """Exit context, popping overrides from thread-local stack.""" + # Pop overrides from the stack + if hasattr(self.config._local, "stack") and self.config._local.stack: # noqa: SLF001 + self.config._local.stack.pop() # noqa: SLF001 + + +# ============================================================================ +# Quantity `__repr__` + +QUANTITY_REPR_CONFIG_KEYS: Final = frozenset( + {"short_arrays", "use_short_name", "named_unit", "indent"} +) + + +class QuantityReprConfig(LocalConfigurable): + """Configuration for quantity ``__repr__()`` display options. + + This controls how quantity objects are displayed in ``repr()``. It is + consumed by ``AbstractQuantity.__repr__`` and so applies to every quantity + class, including the default ``Quantity`` (and ``ParametricQuantity`` from + the ``unxts.parametric`` package). + + Attributes + ---------- + short_arrays : bool | Literal["compact"] + Controls how arrays are displayed in repr. Options: + - "compact": Show array values without Array wrapper + - `True`: Show short array summary (shape/dtype) + - `False`: Show full array representation + Default: `False` + use_short_name : bool + If True and a class has a `short_name` attribute, use the short + name instead of the full class name in repr. Default: False + named_unit : bool + If `True`, display unit as a named argument `unit='m'`. + If `False`, display unit as a positional argument `'m'`. + Default: `True` + + Examples + -------- + >>> import unxt as u + >>> with u.config.quantity_repr.override( + ... short_arrays="compact", use_short_name=True + ... ): + ... q = u.Q([1, 2, 3], "m") + ... print(repr(q)) + Q([1, 2, 3], unit='m') + + """ + + _config_keys: ClassVar[frozenset[str]] = QUANTITY_REPR_CONFIG_KEYS + + short_arrays: ClassVar[object] = Union( + [Bool(), Enum(("compact",))], + default_value=False, + help=( + "Controls array display in repr. " + "Options: 'compact' (values only), " + "True (shape/dtype), False (full repr)" + ), + ).tag(config=True) + + use_short_name: ClassVar[object] = Bool( + default_value=False, + help=( + "Use a class's short name if available (e.g. Quantity -> Q, " + "ParametricQuantity -> PQ)" + ), + ).tag(config=True) + + named_unit: ClassVar[object] = Bool( + default_value=True, + help="Display unit as named argument (unit='m') vs positional ('m')", + ).tag(config=True) + + indent: ClassVar[object] = Int( + default_value=4, help="Indentation level for nested structures in repr" + ).tag(config=True) + + # ============================================================================ # Quantity `__str__` @@ -327,6 +336,8 @@ class QuantityStrConfig(LocalConfigurable): """ + _config_keys: ClassVar[frozenset[str]] = QUANTITY_STR_CONFIG_KEYS + short_arrays: ClassVar[object] = Union( [Bool(), Enum(("compact",))], default_value="compact", @@ -354,101 +365,6 @@ class QuantityStrConfig(LocalConfigurable): default_value=4, help="Indentation level for nested structures in str" ).tag(config=True) - def __getattribute__(self, name: str) -> Any: - """Get attribute, checking thread-local overrides first.""" - if name in QUANTITY_STR_CONFIG_KEYS: - with contextlib.suppress(AttributeError): - local = object.__getattribute__(self, "_local") - if hasattr(local, "stack") and local.stack: - for overrides in reversed(local.stack): - if name in overrides: - return overrides[name] - - return object.__getattribute__(self, name) - - def override( - self, cfg: Config | None = None, /, **kwargs: Any - ) -> "_NestedConfigContext": - """Create a context manager for temporary config changes. - - Parameters - ---------- - cfg : traitlets.config.Config, optional - A traitlets Config object with settings for this config class. - Cannot be used together with keyword arguments. - **kwargs - Configuration options to set temporarily. - Cannot be used together with cfg parameter. - - Returns - ------- - _NestedConfigContext - A context manager that applies the config changes on entry - and restores previous values on exit. - - Raises - ------ - ValueError - If both cfg and keyword arguments are provided. - - """ - if cfg is not None and kwargs: - msg = "Cannot specify both cfg and keyword arguments to override()" - raise ValueError(msg) - - if kwargs: - unknown_keys = set(kwargs) - QUANTITY_STR_CONFIG_KEYS - if unknown_keys: - valid_keys = ", ".join(sorted(QUANTITY_STR_CONFIG_KEYS)) - unknown = ", ".join(sorted(unknown_keys)) - msg = ( - f"Unknown QuantityStrConfig override option(s): {unknown}. " - f"Valid options are: {valid_keys}" - ) - raise ValueError(msg) - - # Validate and resolve values through traitlets immediately so - # override() fails fast with clear errors. - temp_instance = self.__class__() - validated_kwargs: dict[str, Any] = {} - for key, value in kwargs.items(): - try: - setattr(temp_instance, key, value) - except TraitError as e: - msg = ( - "Invalid value for QuantityStrConfig override option " - f"'{key}': {value!r}" - ) - raise ValueError(msg) from e - # Bypass thread-local override lookup when reading back the - # resolved trait value from the temporary instance. - validated_kwargs[key] = object.__getattribute__(temp_instance, key) - kwargs = validated_kwargs - - if cfg is not None: - # Create a temporary instance, apply the config to it, and read the - # resolved trait values. This ensures LazyConfigValue objects are - # properly resolved to their actual values. - temp_instance = self.__class__(config=cfg) - # Only override the traits the Config actually specifies. Iterating - # every trait would read the class default for traits absent from - # ``cfg`` and clobber any globally-configured value on entry. - specified: set[str] = set() - for klass in type(self).__mro__: - section = cfg.get(klass.__name__, None) - if section: - specified |= set(section) - # Intersect with the *override allowlist*, not all traits: - # ``__getattribute__`` only consults these keys, so an entry outside - # it could never be observed -- it would just be a dead push. - overrides = { - name: object.__getattribute__(temp_instance, name) - for name in specified & QUANTITY_STR_CONFIG_KEYS - } - kwargs = overrides - - return _NestedConfigContext(self, kwargs) - # ============================================================================ # Unxt configuration