diff --git a/CHANGES.md b/CHANGES.md index ea82ab3..f4f245e 100644 --- a/CHANGES.md +++ b/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 diff --git a/examples/dynamic.ipynb b/examples/dynamic.ipynb index 95967a5..47c3b46 100644 --- a/examples/dynamic.ipynb +++ b/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, diff --git a/test/data/paramtest.ipynb b/test/data/paramtest.ipynb new file mode 100644 index 0000000..1daf501 --- /dev/null +++ b/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 +} diff --git a/test/test_cli.py b/test/test_cli.py index 253cc63..592448a 100644 --- a/test/test_cli.py +++ b/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( diff --git a/test/test_core.py b/test/test_core.py index 0606d0c..0023daf 100644 --- a/test/test_core.py +++ b/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"), + } diff --git a/test/test_parameters.py b/test/test_parameters.py index d67ed40..dcf0fc6 100644 --- a/test/test_parameters.py +++ b/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": { diff --git a/test/test_wrapper.py b/test/test_wrapper.py index 163d945..9a7757a 100644 --- a/test/test_wrapper.py +++ b/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() diff --git a/xcengine/cli.py b/xcengine/cli.py index ef7f91a..b621c00 100644 --- a/xcengine/cli.py +++ b/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", diff --git a/xcengine/core.py b/xcengine/core.py index fef4161..49a26a0 100755 --- a/xcengine/core.py +++ b/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, diff --git a/xcengine/parameters.py b/xcengine/parameters.py index 5a19c75..a321389 100644 --- a/xcengine/parameters.py +++ b/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]: