Improve automated parameter extraction

CHANGES.md

@@ -1,5 +1,7 @@
 ## Changes in 0.1.1 (in development)
 
+* Improve automated parameter extraction (#9)
+
 ## Changes in 0.1.0
 
 * Initial release

examples/dynamic.ipynb

@@ -212,7 +212,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.5"
+   "version": "3.12.8"
   }
  },
  "nbformat": 4,

test/data/paramtest.ipynb

@@ -0,0 +1,62 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "39da67c4-94a7-4f13-8bda-87a9a3812694",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "\n",
+    "my_constant = math.floor(3.5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bae7e3db-0886-4d33-b982-4374452f62d0",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": [
+     "parameters"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "parameter_1 = my_constant * 2\n",
+    "parameter_2 = \"default value\""
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

test/test_cli.py

@@ -74,14 +74,12 @@ def test_image_build(builder_mock, tmp_path):
     )
     instance_mock.build.assert_called_once_with()
 
+
 @patch("xcengine.cli.ContainerRunner")
 def test_image_run(runner_mock):
     cli_runner = CliRunner()
     instance_mock = runner_mock.return_value = MagicMock()
-    result = cli_runner.invoke(
-        cli,
-        ["image", "run", "foo"]
-    )
+    result = cli_runner.invoke(cli, ["image", "run", "foo"])
     runner_mock.assert_called_once_with(image="foo", output_dir=None)
     assert result.exit_code == 0
     instance_mock.run.assert_called_once_with(

test/test_core.py

@@ -14,7 +14,7 @@
 
 from unittest.mock import MagicMock, patch
 
-from xcengine.core import ChunkStream, ImageBuilder
+from xcengine.core import ChunkStream, ImageBuilder, ScriptCreator
 
 
 @patch("xcengine.core.ScriptCreator.__init__")
@@ -47,13 +47,13 @@ def test_image_builder_init(init_mock, tmp_path, tag):
     init_mock.assert_called_once_with(nb_path)
 
 
-def test_init_runner_invalid_image_type():
+def test_runner_init_invalid_image_type():
     with pytest.raises(ValueError, match='Invalid type "int"'):
         # noinspection PyTypeChecker
         xcengine.core.ContainerRunner(666, pathlib.Path("/foo"))
 
 
-def test_init_runner_with_string():
+def test_runner_init_with_string():
     image_name = "foo"
     image_mock = Mock(docker.models.images.Image)
     client_mock = Mock(docker.client.DockerClient)
@@ -69,7 +69,7 @@ def get_mock(name):
     assert image_mock == runner.image
 
 
-def test_init_runner_with_image():
+def test_runner_init_with_image():
     runner = xcengine.core.ContainerRunner(
         image := Mock(docker.models.images.Image), pathlib.Path("/foo")
     )
@@ -117,3 +117,13 @@ def test_chunk_stream():
     chunk_stream = ChunkStream(bytegen)
     assert chunk_stream.readable()
     assert BufferedReader(chunk_stream).read() == expected
+
+
+def test_script_creator_init():
+    script_creator = ScriptCreator(
+        pathlib.Path(__file__).parent / "data" / "paramtest.ipynb"
+    )
+    assert script_creator.nb_params.params == {
+        "parameter_1": (int, 6),
+        "parameter_2": (str, "default value"),
+    }

test/test_parameters.py

@@ -156,6 +156,25 @@ def test_parameters_from_code(expected_vars):
     )
 
 
+def test_parameters_from_code_with_setup(expected_vars):
+    assert (
+        xcengine.parameters.NotebookParameters.from_code(
+            """
+some_int = 2 * half_of_some_int
+some_float = 3.14159
+some_string = some_uppercase_string.lower()
+some_bool = not not_some_bool
+    """,
+            setup_code="""
+half_of_some_int = 21
+some_uppercase_string = "FOO"
+not_some_bool = True
+            """,
+        ).params
+        == expected_vars
+    )
+
+
 def test_parameters_get_workflow_inputs(notebook_parameters):
     assert notebook_parameters.get_cwl_workflow_inputs() == {
         "some_int": {

test/test_wrapper.py

@@ -6,10 +6,12 @@
 @patch("sys.argv", ["wrapper.py", "--verbose"])
 def test_wrapper(tmp_path, monkeypatch):
     import xcengine
+
     for path in xcengine.__path__:
         monkeypatch.syspath_prepend(path)
-    user_code_path = (tmp_path / "user_code.py")
+    user_code_path = tmp_path / "user_code.py"
     user_code_path.touch()
     os.environ["XC_USER_CODE_PATH"] = str(user_code_path)
     from xcengine import wrapper
+
     xcengine.wrapper.main()

xcengine/cli.py

@@ -31,7 +31,7 @@ def cli(verbose):
     "--from-saved",
     is_flag=True,
     help="If --batch and --server both used, serve datasets from saved Zarrs "
-         "rather than computing them on the fly.",
+    "rather than computing them on the fly.",
 )
 
 notebook_argument = click.argument(
@@ -160,7 +160,7 @@ def build(
     "--batch",
     is_flag=True,
     help="Run the compute engine as a batch script. Use with the --output "
-         "option to copy output out of the container.",
+    "option to copy output out of the container.",
 )
 @click.option(
     "-s",

xcengine/core.py

@@ -2,8 +2,10 @@
 # Permissions are hereby granted under the terms of the MIT License:
 # https://opensource.org/licenses/MIT.
 
+import functools
 import io
 import json
+import operator
 import os
 import shutil
 import sys
@@ -71,8 +73,15 @@ def process_params_cell(self) -> None:
                 params_cell_index = i
                 break
         if params_cell_index is not None:
+            setup_code = "\n".join(
+                map(
+                    operator.attrgetter("source"),
+                    self.notebook.cells[:params_cell_index],
+                )
+            )
             self.nb_params = NotebookParameters.from_code(
-                self.notebook.cells[params_cell_index].source
+                self.notebook.cells[params_cell_index].source,
+                setup_code=setup_code,
             )
             self.notebook.cells.insert(
                 params_cell_index + 1,

xcengine/parameters.py

@@ -35,12 +35,14 @@ def make_cwl_params(self):
             self.cwl_params["product"] = "Directory", None
 
     @classmethod
-    def from_code(cls, code: str) -> "NotebookParameters":
+    def from_code(
+        cls, code: str, setup_code: str | None = None
+    ) -> "NotebookParameters":
         # TODO run whole notebook up to params cell, not just the params cell!
         # (Because it might use imports etc. from earlier in the notebook.)
         # This will need some tweaking of the parameter extraction -- see
         # comment therein.
-        return cls(cls.extract_variables(code))
+        return cls(cls.extract_variables(code, setup_code))
 
     @classmethod
     def from_yaml(cls, yaml_content: str | typing.IO) -> "NotebookParameters":
@@ -61,12 +63,18 @@ def from_yaml_file(cls, path: str | os.PathLike) -> "NotebookParameters":
             return cls.from_yaml(fh)
 
     @classmethod
-    def extract_variables(cls, code: str) -> dict[str, tuple[type, Any]]:
-        # TODO: just working on a single code block is insufficient:
-        # We should execute everything up to the params cell, but only extract
-        # variables defined in the params cell.
-        exec(code, globals(), locals_ := {})
-        return {k: cls.make_param_tuple(locals_[k]) for k in (locals_.keys())}
+    def extract_variables(
+        cls, code: str, setup_code: str | None = None
+    ) -> dict[str, tuple[type, Any]]:
+        if setup_code is None:
+            locals_ = {}
+            old_locals = {}
+        else:
+            exec(setup_code, globals(), locals_ := {})
+            old_locals = locals_.copy()
+        exec(code, globals(), locals_)
+        new_vars = locals_.keys() - old_locals.keys()
+        return {k: cls.make_param_tuple(locals_[k]) for k in new_vars}
 
     @staticmethod
     def make_param_tuple(value: Any) -> tuple[type, Any]: