Add init argument builtins to CodeInput (#91)

This allows one to update the builtins variables when compiling the code
which is necessary when the teacher wants to provide variables to the
code that do not need to be created by the student.

Author: agoscinski

src/scwidgets/code/_widget_code_input.py

@@ -1,4 +1,5 @@
 import ast
+import copy
 import inspect
 import re
 import sys
@@ -7,7 +8,7 @@
 import types
 import warnings
 from functools import wraps
-from typing import List, Optional, Tuple
+from typing import Any, List, Optional, Tuple
 
 from widget_code_input import WidgetCodeInput
 from widget_code_input.utils import (
@@ -34,6 +35,8 @@ class CodeInput(WidgetCodeInput):
         `"x, y = 5"`
     :param docstring: The docstring of the function
     :param function_body: The function definition without indentation
+    :param builtins: A dict of variable name and value that is added to the
+      globals __builtins__ and thus available on initialization
     """
 
     valid_code_themes = ["nord", "solarizedLight", "basicLight"]
@@ -45,6 +48,7 @@ def __init__(
         function_parameters: Optional[str] = None,
         docstring: Optional[str] = None,
         function_body: Optional[str] = None,
+        builtins: Optional[dict[str, Any]] = None,
         code_theme: str = "basicLight",
     ):
         if function is not None:
@@ -69,6 +73,7 @@ def __init__(
         function_parameters = "" if function_parameters is None else function_parameters
         docstring = "\n" if docstring is None else docstring
         function_body = "" if function_body is None else function_body
+        self._builtins = {} if builtins is None else builtins
         super().__init__(
             function_name, function_parameters, docstring, function_body, code_theme
         )
@@ -94,13 +99,17 @@ def unwrapped_function(self) -> types.FunctionType:
         :raise SyntaxError: if the function code has syntax errors (or if
           the function name is not a valid identifier)
         """
+        # we shallow copy the builtins to be able to overwrite it
+        # if self.builtins changes
         globals_dict = {
-            "__builtins__": globals()["__builtins__"],
+            "__builtins__": copy.copy(globals()["__builtins__"]),
             "__name__": "__main__",
             "__doc__": None,
             "__package__": None,
         }
 
+        globals_dict["__builtins__"].update(self._builtins)
+
         if not is_valid_variable_name(self.function_name):
             raise SyntaxError("Invalid function name '{}'".format(self.function_name))
 
@@ -275,6 +284,14 @@ def wrapper(*args, **kwargs):
 
         return catch_exceptions(self.unwrapped_function)
 
+    @property
+    def builtins(self) -> dict[str, Any]:
+        return self._builtins
+
+    @builtins.setter
+    def builtins(self, value: dict[str, Any]):
+        self._builtins = value
+
 
 # Temporary fix until https://github.com/osscar-org/widget-code-input/pull/26
 # is merged

tests/test_code.py

@@ -115,6 +115,25 @@ def test_call(self):
         assert code(1, 1) == 2
         assert code(0, 1) == 1
 
+    def test_builtins(self):
+        """Tests if import work when they are specified by builtins."""
+        import numpy as np
+
+        # to check if annotation works
+        def function(arr: np.ndarray):
+            return arr + builtins_variable  # noqa: F821
+
+        code_input = CodeInput(function, builtins={"np": np, "builtins_variable": 0})
+        code_input.unwrapped_function(np.array([0]))
+
+        # check if builtins is overwritten,
+        # the builtins_variable should not be there anymore afterwards
+        code_input.builtins = {"np": np}
+        with pytest.raises(
+            NameError, match=r".*name 'builtins_variable' is not defined.*"
+        ):
+            code_input.unwrapped_function(np.array([0]))
+
 
 def get_code_exercise(
     checks: List[Check],