path: root/pint/facets/context/registry.py

"""
    pint.facets.context.registry
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    :copyright: 2022 by Pint Authors, see AUTHORS for more details.
    :license: BSD, see LICENSE for more details.
"""

from __future__ import annotations

import functools
from collections import ChainMap
from contextlib import contextmanager
from typing import Any, Callable, ContextManager, Dict, Union

from ..._typing import F
from ...errors import UndefinedUnitError
from ...util import find_connected_nodes, find_shortest_path, logger
from ..plain import PlainRegistry, UnitDefinition
from .definitions import ContextDefinition
from .objects import Context, ContextChain

# TODO: Put back annotation when possible
# registry_cache: "RegistryCache"


class ContextCacheOverlay:
    """Layer on top of the plain UnitRegistry cache, specific to a combination of
    active contexts which contain unit redefinitions.
    """

    def __init__(self, registry_cache) -> None:
        self.dimensional_equivalents = registry_cache.dimensional_equivalents
        self.root_units = {}
        self.dimensionality = registry_cache.dimensionality
        self.parse_unit = registry_cache.parse_unit


class ContextRegistry(PlainRegistry):
    """Handle of Contexts.

    Conversion between units with different dimensions according
    to previously established relations (contexts).
    (e.g. in the spectroscopy, conversion between frequency and energy is possible)

    Capabilities:

    - Register contexts.
    - Enable and disable contexts.
    - Parse @context directive.
    """

    Context = Context

    def __init__(self, **kwargs: Any) -> None:
        # Map context name (string) or abbreviation to context.
        self._contexts: Dict[str, Context] = {}
        # Stores active contexts.
        self._active_ctx = ContextChain()
        # Map context chain to cache
        self._caches = {}
        # Map context chain to units override
        self._context_units = {}

        super().__init__(**kwargs)

        # Allow contexts to add override layers to the units
        self._units = ChainMap(self._units)

    def _register_definition_adders(self) -> None:
        super()._register_definition_adders()
        self._register_adder(ContextDefinition, self.add_context)

    def add_context(self, context: Union[Context, ContextDefinition]) -> None:
        """Add a context object to the registry.

        The context will be accessible by its name and aliases.

        Notice that this method will NOT enable the context;
        see :meth:`enable_contexts`.
        """
        if isinstance(context, ContextDefinition):
            context = Context.from_definition(context, self.get_dimensionality)

        if not context.name:
            raise ValueError("Can't add unnamed context to registry")
        if context.name in self._contexts:
            logger.warning(
                "The name %s was already registered for another context.", context.name
            )
        self._contexts[context.name] = context
        for alias in context.aliases:
            if alias in self._contexts:
                logger.warning(
                    "The name %s was already registered for another context",
                    context.name,
                )
            self._contexts[alias] = context

    def remove_context(self, name_or_alias: str) -> Context:
        """Remove a context from the registry and return it.

        Notice that this methods will not disable the context;
        see :meth:`disable_contexts`.
        """
        context = self._contexts[name_or_alias]

        del self._contexts[context.name]
        for alias in context.aliases:
            del self._contexts[alias]

        return context

    def _build_cache(self, loaded_files=None) -> None:
        super()._build_cache(loaded_files)
        self._caches[()] = self._cache

    def _switch_context_cache_and_units(self) -> None:
        """If any of the active contexts redefine units, create variant self._cache
        and self._units specific to the combination of active contexts.
        The next time this method is invoked with the same combination of contexts,
        reuse the same variant self._cache and self._units as in the previous time.
        """
        del self._units.maps[:-1]
        units_overlay = any(ctx.redefinitions for ctx in self._active_ctx.contexts)
        if not units_overlay:
            # Use the default _cache and _units
            self._cache = self._caches[()]
            return

        key = self._active_ctx.hashable()
        try:
            self._cache = self._caches[key]
            self._units.maps.insert(0, self._context_units[key])
        except KeyError:
            pass

        # First time using this specific combination of contexts and it contains
        # unit redefinitions
        base_cache = self._caches[()]
        self._caches[key] = self._cache = ContextCacheOverlay(base_cache)

        self._context_units[key] = units_overlay = {}
        self._units.maps.insert(0, units_overlay)

        on_redefinition_backup = self._on_redefinition
        self._on_redefinition = "ignore"
        try:
            for ctx in reversed(self._active_ctx.contexts):
                for definition in ctx.redefinitions:
                    self._redefine(definition)
        finally:
            self._on_redefinition = on_redefinition_backup

    def _redefine(self, definition: UnitDefinition) -> None:
        """Redefine a unit from a context"""
        # Find original definition in the UnitRegistry
        candidates = self.parse_unit_name(definition.name)
        if not candidates:
            raise UndefinedUnitError(definition.name)
        candidates_no_prefix = [c for c in candidates if not c[0]]
        if not candidates_no_prefix:
            raise ValueError(f"Can't redefine a unit with a prefix: {definition.name}")
        assert len(candidates_no_prefix) == 1
        _, name, _ = candidates_no_prefix[0]
        try:
            basedef = self._units[name]
        except KeyError:
            raise UndefinedUnitError(name)

        # Rebuild definition as a variant of the plain
        if basedef.is_base:
            raise ValueError("Can't redefine a plain unit to a derived one")

        dims_old = self._get_dimensionality(basedef.reference)
        dims_new = self._get_dimensionality(definition.reference)
        if dims_old != dims_new:
            raise ValueError(
                f"Can't change dimensionality of {basedef.name} "
                f"from {dims_old} to {dims_new} in a context"
            )

        # Do not modify in place the original definition, as (1) the context may
        # be shared by other registries, and (2) it would alter the cache key
        definition = UnitDefinition(
            name=basedef.name,
            defined_symbol=basedef.symbol,
            aliases=basedef.aliases,
            reference=definition.reference,
            converter=definition.converter,
        )

        # Write into the context-specific self._units.maps[0] and self._cache.root_units
        self.define(definition)

    def enable_contexts(
        self, *names_or_contexts: Union[str, Context], **kwargs
    ) -> None:
        """Enable contexts provided by name or by object.

        Parameters
        ----------
        *names_or_contexts :
            one or more contexts or context names/aliases
        **kwargs :
            keyword arguments for the context(s)

        Examples
        --------
        See :meth:`context`
        """

        # If present, copy the defaults from the containing contexts
        if self._active_ctx.defaults:
            kwargs = dict(self._active_ctx.defaults, **kwargs)

        # For each name, we first find the corresponding context
        ctxs = [
            self._contexts[name] if isinstance(name, str) else name
            for name in names_or_contexts
        ]

        # Check if the contexts have been checked first, if not we make sure
        # that dimensions are expressed in terms of plain dimensions.
        for ctx in ctxs:
            if ctx.checked:
                continue
            funcs_copy = dict(ctx.funcs)
            for (src, dst), func in funcs_copy.items():
                src_ = self._get_dimensionality(src)
                dst_ = self._get_dimensionality(dst)
                if src != src_ or dst != dst_:
                    ctx.remove_transformation(src, dst)
                    ctx.add_transformation(src_, dst_, func)
            ctx.checked = True

        # and create a new one with the new defaults.
        contexts = tuple(Context.from_context(ctx, **kwargs) for ctx in ctxs)

        # Finally we add them to the active context.
        self._active_ctx.insert_contexts(*contexts)
        self._switch_context_cache_and_units()

    def disable_contexts(self, n: int = None) -> None:
        """Disable the last n enabled contexts.

        Parameters
        ----------
        n : int
            Number of contexts to disable. Default: disable all contexts.
        """
        self._active_ctx.remove_contexts(n)
        self._switch_context_cache_and_units()

    @contextmanager
    def context(self, *names, **kwargs) -> ContextManager[Context]:
        """Used as a context manager, this function enables to activate a context
        which is removed after usage.

        Parameters
        ----------
        *names : name(s) of the context(s).
        **kwargs : keyword arguments for the contexts.

        Examples
        --------
        Context can be called by their name:

        >>> import pint.facets.context.objects
        >>> import pint
        >>> ureg = pint.UnitRegistry()
        >>> ureg.add_context(pint.facets.context.objects.Context('one'))
        >>> ureg.add_context(pint.facets.context.objects.Context('two'))
        >>> with ureg.context('one'):
        ...     pass

        If a context has an argument, you can specify its value as a keyword argument:

        >>> with ureg.context('one', n=1):
        ...     pass

        Multiple contexts can be entered in single call:

        >>> with ureg.context('one', 'two', n=1):
        ...     pass

        Or nested allowing you to give different values to the same keyword argument:

        >>> with ureg.context('one', n=1):
        ...     with ureg.context('two', n=2):
        ...         pass

        A nested context inherits the defaults from the containing context:

        >>> with ureg.context('one', n=1):
        ...     # Here n takes the value of the outer context
        ...     with ureg.context('two'):
        ...         pass
        """
        # Enable the contexts.
        self.enable_contexts(*names, **kwargs)

        try:
            # After adding the context and rebuilding the graph, the registry
            # is ready to use.
            yield self
        finally:
            # Upon leaving the with statement,
            # the added contexts are removed from the active one.
            self.disable_contexts(len(names))

    def with_context(self, name, **kwargs) -> Callable[[F], F]:
        """Decorator to wrap a function call in a Pint context.

        Use it to ensure that a certain context is active when
        calling a function.

        Parameters
        ----------
        name :
            name of the context.
        **kwargs :
            keyword arguments for the context


        Returns
        -------
        callable: the wrapped function.

        Examples
        --------
        >>> @ureg.with_context('sp')
        ... def my_cool_fun(wavelength):
        ...     print('This wavelength is equivalent to: %s', wavelength.to('terahertz'))
        """

        def decorator(func):
            assigned = tuple(
                attr for attr in functools.WRAPPER_ASSIGNMENTS if hasattr(func, attr)
            )
            updated = tuple(
                attr for attr in functools.WRAPPER_UPDATES if hasattr(func, attr)
            )

            @functools.wraps(func, assigned=assigned, updated=updated)
            def wrapper(*values, **wrapper_kwargs):
                with self.context(name, **kwargs):
                    return func(*values, **wrapper_kwargs)

            return wrapper

        return decorator

    def _convert(self, value, src, dst, inplace=False):
        """Convert value from some source to destination units.

        In addition to what is done by the PlainRegistry,
        converts between units with different dimensions by following
        transformation rules defined in the context.

        Parameters
        ----------
        value :
            value
        src : UnitsContainer
            source units.
        dst : UnitsContainer
            destination units.
        inplace :
             (Default value = False)

        Returns
        -------
        callable
            converted value
        """
        # If there is an active context, we look for a path connecting source and
        # destination dimensionality. If it exists, we transform the source value
        # by applying sequentially each transformation of the path.
        if self._active_ctx:

            src_dim = self._get_dimensionality(src)
            dst_dim = self._get_dimensionality(dst)

            path = find_shortest_path(self._active_ctx.graph, src_dim, dst_dim)
            if path:
                src = self.Quantity(value, src)
                for a, b in zip(path[:-1], path[1:]):
                    src = self._active_ctx.transform(a, b, self, src)

                value, src = src._magnitude, src._units

        return super()._convert(value, src, dst, inplace)

    def _get_compatible_units(self, input_units, group_or_system):
        src_dim = self._get_dimensionality(input_units)

        ret = super()._get_compatible_units(input_units, group_or_system)

        if self._active_ctx:
            ret = ret.copy()  # Do not alter self._cache
            nodes = find_connected_nodes(self._active_ctx.graph, src_dim)
            if nodes:
                for node in nodes:
                    ret |= self._cache.dimensional_equivalents[node]

        return ret