masque/masque/library/lazy.py

"""Closure-backed lazy library implementation."""
from __future__ import annotations

from pprint import pformat
from typing import TYPE_CHECKING, Self, cast
import logging

from ..error import LibraryError
from .base import ILibrary
from .capabilities import IMaterializable

if TYPE_CHECKING:
    from collections.abc import Callable, Iterator, Mapping, Sequence

    from ..pattern import Pattern


logger = logging.getLogger(__name__)


class LazyLibrary(ILibrary, IMaterializable):
    """
    This class is usually used to create a library of Patterns by mapping names to
     functions which generate or load the relevant `Pattern` object as-needed.

    TODO: lots of stuff causes recursive loads (e.g. data_to_ports?). What should you avoid?
    """
    mapping: dict[str, Callable[[], Pattern]]
    cache: dict[str, Pattern]
    _lookups_in_progress: list[str]

    def __init__(self) -> None:
        self.mapping = {}
        self.cache = {}
        self._lookups_in_progress = []

    def __setitem__(
            self,
            key: str,
            value: Pattern | Callable[[], Pattern],
            ) -> None:
        if key in self.mapping:
            raise LibraryError(f'"{key}" already exists in the library. Overwriting is not allowed!')

        if callable(value):
            value_func = value
        else:
            value_func = lambda: cast('Pattern', value)      # noqa: E731

        self.mapping[key] = value_func
        if key in self.cache:
            del self.cache[key]

    def __delitem__(self, key: str) -> None:
        del self.mapping[key]
        if key in self.cache:
            del self.cache[key]

    def __getitem__(self, key: str) -> Pattern:
        return self.materialize(key, persist=True)

    def materialize(self, key: str, *, persist: bool = True) -> Pattern:
        logger.debug(f'loading {key}')
        if key in self.cache:
            logger.debug(f'found {key} in cache')
            return self.cache[key]

        if key in self._lookups_in_progress:
            chain = ' -> '.join(self._lookups_in_progress + [key])
            raise LibraryError(
                f'Detected circular reference or recursive lookup of "{key}".\n'
                f'Lookup chain: {chain}\n'
                'This may be caused by an invalid (cyclical) reference, or buggy code.\n'
                'If you are lazy-loading a file, try a non-lazy load and check for reference cycles.'
                )

        self._lookups_in_progress.append(key)
        try:
            func = self.mapping[key]
            pat = func()
        finally:
            self._lookups_in_progress.pop()
        if persist:
            self.cache[key] = pat
        return pat

    def __iter__(self) -> Iterator[str]:
        return iter(self.mapping)

    def __len__(self) -> int:
        return len(self.mapping)

    def __contains__(self, key: object) -> bool:
        return key in self.mapping

    def referenced_patterns(
            self,
            tops: str | Sequence[str] | None = None,
            skip: set[str] | None = None,
            ) -> set[str]:
        # Closure-backed cells do not have hierarchy metadata. Preserve laziness
        # by loading only patterns reached from the requested roots.
        return self._referenced_patterns_by_lookup(tops=tops, skip=skip)

    def _merge(self, key_self: str, other: Mapping[str, Pattern], key_other: str) -> None:
        if isinstance(other, LazyLibrary):
            self.mapping[key_self] = other.mapping[key_other]
            if key_other in other.cache:
                self.cache[key_self] = other.cache[key_other]
        else:
            self[key_self] = other[key_other]

    def __repr__(self) -> str:
        return '<LazyLibrary with keys\n' + pformat(list(self.keys())) + '>'

    def rename(
            self,
            old_name: str,
            new_name: str,
            move_references: bool = False,
            ) -> Self:
        """
        Rename a pattern.

        Args:
            old_name: Current name for the pattern
            new_name: New name for the pattern
            move_references: Whether to scan all refs in the pattern and
                move them to point to `new_name` as necessary.
                Default `False`.

        Returns:
            self
        """
        if old_name not in self.mapping:
            raise LibraryError(f'"{old_name}" does not exist in the library.')
        if old_name == new_name:
            return self

        self[new_name] = self.mapping[old_name]        # copy over function
        if old_name in self.cache:
            self.cache[new_name] = self.cache[old_name]
        del self[old_name]

        if move_references:
            self.move_references(old_name, new_name)

        return self

    def move_references(self, old_target: str, new_target: str) -> Self:
        """
        Change all references pointing at `old_target` into references pointing at `new_target`.

        Args:
            old_target: Current reference target
            new_target: New target for the reference

        Returns:
            self
        """
        if old_target == new_target:
            return self

        self.precache()
        for pattern in self.cache.values():
            if old_target in pattern.refs:
                pattern.refs[new_target].extend(pattern.refs[old_target])
                del pattern.refs[old_target]
        return self

    def precache(self) -> Self:
        """
        Force all patterns into the cache

        Returns:
            self
        """
        for key in self.mapping:
            _ = self[key]       # want to trigger our own __getitem__
        return self

    def __deepcopy__(self, memo: dict | None = None) -> LazyLibrary:
        raise LibraryError('LazyLibrary cannot be deepcopied (deepcopy doesn\'t descend into closures)')