context.py

"""
Provides classes for tracing the execution of operations.

The operations include key generation, scalar multiplication, formula execution and individual operation evaluation.
These operations are traced in `Context` classes using `Actions`. Different contexts trace actions differently.

A :py:class:`DefaultContext` traces actions into a tree as they are executed (a scalar
multiplication actions has as its children an ordered list of the individual formula executions it has done).

A :py:class:`PathContext` works like a :py:class:`DefaultContext` that only traces an action on a particular path
in the tree.
"""
from abc import abstractmethod, ABC
from collections import OrderedDict
from copy import deepcopy
from typing import List, Optional, ContextManager, Any, Tuple, Sequence, Callable

from public import public


@public
class Action:
    """An Action."""

    inside: bool

    def __init__(self):
        self.inside = False

    def __enter__(self):
        if current is not None:
            current.enter_action(self)
        self.inside = True
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if current is not None:
            current.exit_action(self)
        self.inside = False


@public
class ResultAction(Action):
    """An action that has a result."""

    _result: Any = None
    _has_result: bool = False

    @property
    def result(self) -> Any:
        if not self._has_result:
            raise AttributeError("No result set")
        return self._result

    def exit(self, result: Any):
        if not self.inside:
            raise RuntimeError("Result set outside of action scope")
        if self._has_result:
            return
        self._has_result = True
        self._result = result
        return result

    def __exit__(self, exc_type, exc_val, exc_tb):
        if (
                not self._has_result
                and exc_type is None
                and exc_val is None
                and exc_tb is None
        ):
            raise RuntimeError("Result unset on action exit")
        super().__exit__(exc_type, exc_val, exc_tb)


@public
class Tree(OrderedDict):
    """A recursively-implemented tree."""

    def get_by_key(self, path: List) -> Any:
        """
        Get the value in the tree at a position given by the path.

        :param path: The path to get.
        :return: The value in the tree.
        """
        if len(path) == 0:
            return self
        value = self[path[0]]
        if len(path) == 1:
            return value
        elif isinstance(value, Tree):
            return value.get_by_key(path[1:])
        else:
            raise ValueError

    def get_by_index(self, path: List[int]) -> Tuple[Any, Any]:
        """
        Get the key and value in the tree at a position given by the path of indices.

        The nodes inside a level of a tree are ordered by insertion order.

        :param path: The path to get.
        :return: The key and value.
        """
        if len(path) == 0:
            raise ValueError
        key = list(self.keys())[path[0]]
        value = self[key]
        if len(path) == 1:
            return key, value
        elif isinstance(value, Tree):
            return value.get_by_index(path[1:])
        else:
            raise ValueError

    def repr(self, depth: int = 0) -> str:
        """
        Construct a textual representation of the tree. Useful for visualization and debugging.

        :param depth:
        :return: The resulting textual representation.
        """
        result = ""
        for key, value in self.items():
            if isinstance(value, Tree):
                result += "\t" * depth + str(key) + "\n"
                result += value.repr(depth + 1)
            else:
                result += "\t" * depth + str(key) + ":" + str(value) + "\n"
        return result

    def walk(self, callback: Callable[[Any], None]) -> None:
        """
        Walk the tree, depth-first, with the callback.

        :param callback: The callback to call for all values in the tree.
        """
        for key, val in self.items():
            callback(key)
            if isinstance(val, Tree):
                val.walk(callback)

    def __repr__(self):
        return self.repr()


@public
class Context(ABC):
    """
    Context is an object that traces actions which happen.

    There is always one context active, see functions :py:func:`getcontext`,
    :py:func:`setcontext` and :py:func:`resetcontext`. Also, the :py:func:`local`
    contextmanager.
    """

    @abstractmethod
    def enter_action(self, action: Action) -> None:
        """
        Enter into an action (i.e. start executing it).

        :param action: The action.
        """
        raise NotImplementedError

    @abstractmethod
    def exit_action(self, action: Action) -> None:
        """
        Exit from an action (i.e. stop executing it).

        :param action: The action.
        """
        raise NotImplementedError

    def __str__(self):
        return self.__class__.__name__


@public
class DefaultContext(Context):
    """Context that traces executions of actions in a tree."""

    actions: Tree
    current: List[Action]

    def __init__(self):
        self.actions = Tree()
        self.current = []

    def enter_action(self, action: Action) -> None:
        self.actions.get_by_key(self.current)[action] = Tree()
        self.current.append(action)

    def exit_action(self, action: Action) -> None:
        if len(self.current) < 1 or self.current[-1] != action:
            raise ValueError
        self.current.pop()

    def __repr__(self):
        return f"{self.__class__.__name__}({self.actions!r}, current={self.current!r})"


@public
class PathContext(Context):
    """Context that traces targeted actions."""

    path: List[int]
    current: List[int]
    current_depth: int
    value: Any

    def __init__(self, path: Sequence[int]):
        """
        Create a :py:class:`PathContext`.

        :param path: The path of an action in the execution tree that will be captured.
        """
        self.path = list(path)
        self.current = []
        self.current_depth = 0
        self.value = None

    def enter_action(self, action: Action) -> None:
        if self.current_depth == len(self.current):
            self.current.append(0)
        else:
            self.current[self.current_depth] += 1
        self.current_depth += 1
        if self.path == self.current[: self.current_depth]:
            self.value = action

    def exit_action(self, action: Action) -> None:
        if self.current_depth != len(self.current):
            self.current.pop()
        self.current_depth -= 1

    def __repr__(self):
        return (
            f"{self.__class__.__name__}({self.current!r}, depth={self.current_depth!r})"
        )


current: Optional[Context] = None


class _ContextManager:
    def __init__(self, new_context):
        # TODO: Is this deepcopy a good idea?
        self.new_context = deepcopy(new_context)

    def __enter__(self) -> Optional[Context]:
        global current  # This is OK, skipcq: PYL-W0603
        self.old_context = current
        current = self.new_context
        return current

    def __exit__(self, t, v, tb):
        global current  # This is OK, skipcq: PYL-W0603
        current = self.old_context


@public
def local(ctx: Optional[Context] = None) -> ContextManager:
    """
    Use a local context.

    Use it like a contextmanager, the context is active during its execution.

    >>> with local(DefaultContext()) as ctx:
    ...     with Action() as action:
    ...         pass
    >>> list(ctx.actions)[0] == action
    True

    :param ctx: If none, current context is copied.
    :return: A context manager.
    """
    return _ContextManager(ctx)