Merge pull request #10 from tillahoffmann/tb

Improve tracebacks.

README.rst

@@ -20,7 +20,7 @@ Interactive python sessions, such as `Jupyter notebooks <https://jupyter.org/>`_
    ...     print(a)
    Traceback (most recent call last):
      ...
-   ValueError: `a` is not a permitted global
+   localscope.LocalscopeException: `a` is not a permitted global (file "...", line 1, in print_a)
 
 Motivation and detailed example
 -------------------------------
@@ -65,7 +65,7 @@ This example may seem contrived. But unintended information leakage from the glo
    ...     return sum(((x - y) / sigma) ** 2 for x, y in zip(xs, ys))
    Traceback (most recent call last):
      ...
-   ValueError: `sigma` is not a permitted global
+   localscope.LocalscopeException: `sigma` is not a permitted global (file "...", line 3, in <genexpr>)
 
 Interface
 ---------

localscope/__init__.py

@@ -16,7 +16,6 @@ def localscope(
     predicate: Optional[Callable] = None,
     allowed: Optional[Set[str]] = None,
     allow_closure: bool = False,
-    _globals: Optional[Dict[str, Any]] = None,
 ):
     """
     Restrict the scope of a callable to local variables to avoid unintentional
@@ -27,8 +26,6 @@ def localscope(
         predicate : Predicate to determine whether a global variable is allowed in the
             scope. Defaults to allow any module.
         allowed: Names of globals that are allowed to enter the scope.
-        _globals : Globals associated with the root callable which are passed to
-            dependent code blocks for analysis.
 
     Attributes:
         mfc: Decorator allowing *m*\\ odules, *f*\\ unctions, and *c*\\ lasses to enter
@@ -44,7 +41,8 @@ def localscope(
         ...     print(a)
         Traceback (most recent call last):
         ...
-        ValueError: `a` is not a permitted global
+        localscope.LocalscopeException: `a` is not a permitted global (file "...",
+            line 1, in print_a)
 
         The scope of a function can be extended by providing a list of allowed
         exceptions.
@@ -85,53 +83,111 @@ def localscope(
         blocks) at the time of declaration because static analysis has a minimal impact
         on performance and it is easier to implement.
     """
-    # Set defaults
-    predicate = predicate or inspect.ismodule
+    # Set defaults and construct partial if the callable has not yet been provided for
+    # parameterized decorators, e.g., @localscope(allowed={"foo", "bar"}). This is a
+    # thin wrapper around the actual implementation `_localscope`. The wrapper
+    # reconstructs an informative traceback.
     allowed = set(allowed) if allowed else set()
-    if func is None:
+    predicate = predicate or inspect.ismodule
+    if not func:
         return ft.partial(
             localscope,
             allow_closure=allow_closure,
-            predicate=predicate,
             allowed=allowed,
+            predicate=predicate,
         )
 
+    return _localscope(
+        func,
+        allow_closure=allow_closure,
+        allowed=allowed,
+        predicate=predicate,
+        _globals={},
+    )
+
+
+class LocalscopeException(RuntimeError):
+    """
+    Raised when a callable tries to access a non-local variable.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        code: types.CodeType,
+        instruction: Optional[dis.Instruction] = None,
+    ) -> None:
+        if instruction and instruction.starts_line:
+            lineno = instruction.starts_line
+        else:
+            lineno = code.co_firstlineno
+        details = f'file "{code.co_filename}", line {lineno}, in {code.co_name}'
+        super().__init__(f"{message} ({details})")
+
+
+def _localscope(
+    func: Union[types.FunctionType, types.CodeType],
+    *,
+    predicate: Callable,
+    allowed: Set[str],
+    allow_closure: bool,
+    _globals: Dict[str, Any],
+):
+    """
+    Args:
+        ...: Same as for the wrapper :func:`localscope`.
+        _globals : Globals associated with the root callable which are passed to
+            dependent code blocks for analysis.
+    """
+
+    # Extract global variables from a function
+    # (https://docs.python.org/3/library/types.html#types.FunctionType) or keep the
+    # explicitly provided globals for code objects
+    # (https://docs.python.org/3/library/types.html#types.CodeType).
     if isinstance(func, types.FunctionType):
         code = func.__code__
         _globals = {**func.__globals__, **inspect.getclosurevars(func).nonlocals}
     else:
         code = func
-        _globals = _globals or {}
 
-    # Add function arguments to the list of allowed exceptions
+    # Add function arguments to the list of allowed exceptions.
     allowed.update(code.co_varnames[: code.co_argcount])
 
-    opnames = {"LOAD_GLOBAL"}
+    # Construct set of forbidden operations. The first accesses global variables. The
+    # second accesses variables from the outer scope.
+    forbidden_opnames = {"LOAD_GLOBAL"}
     if not allow_closure:
-        opnames.add("LOAD_DEREF")
+        forbidden_opnames.add("LOAD_DEREF")
 
     LOGGER.info("analysing instructions for %s...", func)
     for instruction in dis.get_instructions(code):
         LOGGER.info(instruction)
         name = instruction.argval
-        if instruction.opname in opnames:
-            # Explicitly allowed
+        if instruction.opname in forbidden_opnames:
+            # Variable explicitly allowed by name or in `builtins`.
             if name in allowed or hasattr(builtins, name):
                 continue
-            # Complain if the variable is not available
+            # Complain if the variable is not available.
             if name not in _globals:
-                raise NameError(f"`{name}` is not in globals")
-            # Get the value of the variable and check it against the predicate
+                raise LocalscopeException(
+                    f"`{name}` is not in globals", code, instruction
+                )
+            # Check if variable is allowed by value.
             value = _globals[name]
             if not predicate(value):
-                raise ValueError(f"`{name}` is not a permitted global")
+                raise LocalscopeException(
+                    f"`{name}` is not a permitted global", code, instruction
+                )
         elif instruction.opname == "STORE_DEREF":
+            # Store a new allowed variable which has been created in the scope of the
+            # function.
             allowed.add(name)
+
     # Deal with code objects recursively after adding the current arguments to the
     # allowed exceptions
     for const in code.co_consts:
         if isinstance(const, types.CodeType):
-            localscope(
+            _localscope(
                 const,
                 _globals=_globals,
                 allow_closure=True,

tests/test_localscope.py

@@ -1,4 +1,4 @@
-from localscope import localscope
+from localscope import localscope, LocalscopeException
 import uuid
 import pytest
 
@@ -16,15 +16,24 @@ def add(a, b):
 
 
 def test_missing_global():
-    with pytest.raises(NameError):
+    def func():
+        return never_declared  # noqa: F821
 
-        @localscope
-        def func():
-            return never_ever_declared  # noqa: F821
+    with pytest.raises(LocalscopeException, match="`never_declared` is not in globals"):
+        localscope(func)
+
+    # IMPORTANT! This function can be executed, but localscope complains because the
+    # global variable is not defined at the time when the function is analysed. This
+    # could be improved, but, most likely, one shouldn't write functions that rely on
+    # future globals in the first place.
+    """
+    never_declared = 123
+    assert func() == 123
+    """
 
 
 def test_forbidden_global():
-    with pytest.raises(ValueError):
+    with pytest.raises(LocalscopeException, match="`forbidden_global` is not a perm"):
 
         @localscope
         def return_forbidden_global():
@@ -57,7 +66,7 @@ def return_forbidden_closure():
 
         return return_forbidden_closure()
 
-    with pytest.raises(ValueError):
+    with pytest.raises(LocalscopeException, match="`forbidden_closure` is not a perm"):
         wrapper()
 
 
@@ -76,7 +85,7 @@ def return_forbidden_closure():
 
 def test_allow_custom_predicate():
     decorator = localscope(predicate=lambda x: isinstance(x, int))
-    with pytest.raises(ValueError):
+    with pytest.raises(LocalscopeException, match="`forbidden_global` is not a perm"):
 
         @decorator
         def return_forbidden_global():
@@ -90,15 +99,15 @@ def return_integer_global():
 
 
 def test_comprehension():
-    with pytest.raises(ValueError):
+    with pytest.raises(LocalscopeException, match="`integer_global` is not a perm"):
 
         @localscope
         def evaluate_mse(xs, ys):  # missing argument integer_global
             return sum(((x - y) / integer_global) ** 2 for x, y in zip(xs, ys))
 
 
 def test_recursive():
-    with pytest.raises(ValueError):
+    with pytest.raises(LocalscopeException, match="`forbidden_global` is not a perm"):
 
         @localscope
         def wrapper():
@@ -108,6 +117,17 @@ def return_forbidden_global():
             return return_forbidden_global()
 
 
+def test_recursive_without_call():
+    # We even raise an exception if we don't call a function. That's necessary because
+    # we can't trace all possible execution paths without actually running the function.
+    with pytest.raises(LocalscopeException, match="`forbidden_global` is not a perm"):
+
+        @localscope
+        def wrapper():
+            def return_forbidden_global():
+                return forbidden_global
+
+
 def test_recursive_local_closure():
     @localscope
     def wrapper():
@@ -134,7 +154,7 @@ def doit():
 
     x = 1
 
-    with pytest.raises(ValueError):
+    with pytest.raises(LocalscopeException, match="`x` is not a permitted"):
 
         @localscope.mfc
         def breakit():