1 files changed, 71 insertions, 17 deletions

diff --git a/pyecsca/ec/context.py b/pyecsca/ec/context.py
index d6f56af..c4e1a3c 100644
--- a/pyecsca/ec/context.py
+++ b/pyecsca/ec/context.py
@@ -1,8 +1,8 @@
-from abc import ABCMeta, abstractmethod
+from abc import ABCMeta, abstractmethod, ABC
 from collections import OrderedDict
 from contextvars import ContextVar, Token
 from copy import deepcopy
-from typing import List, Optional, ContextManager, Any
+from typing import List, Optional, ContextManager, Any, Tuple
 
 from public import public
 
@@ -25,33 +25,87 @@ class Action(object):
         self.inside = False
 
 
-
-
-
+@public
 class Tree(OrderedDict):
 
-    def walk_get(self, path: List) -> Any:
+    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 isinstance(value, Tree):
-            return value.walk_get(path[1:])
+            return value.get_by_key(path[1:])
         elif len(path) == 1:
             return value
         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 __repr__(self):
+        return self.repr()
+
 
 @public
-class Context(object):
-    __metaclass__ = ABCMeta
+class Context(ABC):
+    """A 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`.
+    """
 
     @abstractmethod
-    def enter_action(self, action: Action):
+    def enter_action(self, action: Action) -> None:
+        """
+        Enter into an action (i.e. start executing it).
+
+        :param action: The action.
+        """
         ...
 
     @abstractmethod
-    def exit_action(self, action: Action):
+    def exit_action(self, action: Action) -> None:
+        """
+        Exit from an action (i.e. stop executing it).
+
+        :param action: The action.
+        """
         ...
 
     def __str__(self):
@@ -62,24 +116,24 @@ class Context(object):
 class NullContext(Context):
     """A context that does not trace any actions."""
 
-    def enter_action(self, action: Action):
+    def enter_action(self, action: Action) -> None:
         pass
 
-    def exit_action(self, action: Action):
+    def exit_action(self, action: Action) -> None:
         pass
 
 
 @public
 class DefaultContext(Context):
-    """A context that traces executions of actions."""
+    """A context that traces executions of actions in a tree."""
     actions: Tree
     current: List[Action]
 
-    def enter_action(self, action: Action):
-        self.actions.walk_get(self.current)[action] = Tree()
+    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):
+    def exit_action(self, action: Action) -> None:
         if self.current[-1] != action:
             raise ValueError
         self.current.pop()