"""
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
[docs]
@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
[docs]
@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
[docs]
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)
[docs]
@public
class Tree(OrderedDict):
"""A recursively-implemented tree."""
[docs]
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
[docs]
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
[docs]
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
[docs]
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()
[docs]
@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`.
"""
[docs]
@abstractmethod
def enter_action(self, action: Action) -> None:
"""
Enter into an action (i.e. start executing it).
:param action: The action.
"""
raise NotImplementedError
[docs]
@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__
[docs]
@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 = []
[docs]
def enter_action(self, action: Action) -> None:
self.actions.get_by_key(self.current)[action] = Tree()
self.current.append(action)
[docs]
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})"
[docs]
@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
[docs]
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
[docs]
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
[docs]
@public
def local(ctx: Optional[Context] = None) -> ContextManager:
"""
Use a local context.
:param ctx: If none, current context is copied.
:return: A context manager.
"""
return _ContextManager(ctx)
