Add script_args parameter to ContainerRunner.run

Related to #57

Author: pont-us

test/test_core.py

@@ -29,7 +29,9 @@
 
 @patch("xcengine.core.ScriptCreator.__init__")
 @pytest.mark.parametrize("tag", [None, "bar"])
-@pytest.mark.parametrize("env_file_name", ["environment.yml", "foo.yaml", None])
+@pytest.mark.parametrize(
+    "env_file_name", ["environment.yml", "foo.yaml", None]
+)
 @pytest.mark.parametrize("use_env_file_param", [False, True])
 def test_image_builder_init(
     init_mock,
@@ -56,7 +58,11 @@ def test_image_builder_init(
     )
     assert ib.notebook == nb_path
     assert ib.build_dir == build_path
-    expected_env = environment_path if (use_env_file_param or env_file_name == "environment.yml") else None
+    expected_env = (
+        environment_path
+        if (use_env_file_param or env_file_name == "environment.yml")
+        else None
+    )
     assert ib.environment == expected_env
     if tag is None:
         assert abs(
@@ -123,6 +129,28 @@ def test_runner_run_keep(keep: bool):
         container.remove.assert_called_once_with(force=True)
 
 
+def test_runner_extra_args():
+    runner = xcengine.core.ContainerRunner(
+        image := Mock(docker.models.images.Image),
+        None,
+        client := Mock(DockerClient),
+    )
+    image.tags = []
+    client.containers.run.return_value = (container := MagicMock(Container))
+    container.status = "exited"
+    script_args = ["--foo", "--bar", "42", "--baz", "somestring"]
+    runner.run(
+        run_batch=False,
+        host_port=None,
+        from_saved=False,
+        keep=False,
+        script_args=script_args,
+    )
+    run_args = client.containers.run.call_args
+    command = run_args[1]["command"]
+    assert command == ["python", "execute.py"] + script_args
+
+
 def test_runner_sigint():
     runner = xcengine.core.ContainerRunner(
         image := Mock(docker.models.images.Image),

xcengine/core.py

@@ -33,6 +33,7 @@
 LOGGER = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
+
 class ScriptCreator:
     """Turn a Jupyter notebook into a set of scripts"""
 
@@ -221,7 +222,7 @@ def __init__(
                 self.environment = notebook.parent / nb_env
             else:
                 LOGGER.info(f"No environment specified in notebook.")
-                LOGGER.info(f"Looking for a file named \"environment.yml\".")
+                LOGGER.info(f'Looking for a file named "environment.yml".')
                 notebook_sibling = notebook.parent / "environment.yml"
                 if notebook_sibling.is_file():
                     self.environment = notebook_sibling
@@ -375,6 +376,7 @@ def run(
         host_port: int | None,
         from_saved: bool,
         keep: bool,
+        script_args: list[str] | None = None,
     ):
         LOGGER.info(f"Running container from image {self.image.short_id}")
         LOGGER.info(f"Image tags: {' '.join(self.image.tags)}")
@@ -391,6 +393,9 @@ def run(
                 else []
             )
             + (["--from-saved"] if from_saved else [])
+            + script_args
+            if script_args is not None
+            else []
         )
         run_args: dict[str, Any] = dict(
             image=self.image, command=command, remove=False, detach=True
@@ -400,10 +405,14 @@ def run(
         container: Container = self.client.containers.run(**run_args)
         LOGGER.info(f"Waiting for container {container.short_id} to complete.")
         default_sigint_handler = signal.getsignal(signal.SIGINT)
+
         def signal_hander(signum, frame):
             signal.signal(signal.SIGINT, default_sigint_handler)
-            LOGGER.info(f"Caught SIGINT. Stopping container {container.short_id}")
+            LOGGER.info(
+                f"Caught SIGINT. Stopping container {container.short_id}"
+            )
             container.stop()
+
         signal.signal(signal.SIGINT, signal_hander)
         while container.status in {"created", "running"}:
             LOGGER.debug(