allow only kwargs with Beaker.experiment.create() (#263)

allenai · Dec 15, 2023 · da3d95a · da3d95a
1 parent 5cdc035
commit da3d95a
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,10 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Fixed
+
+- Allow only key-word arguments with `Beaker.experiment.create()`.
+
 ## [v1.22.0](https://github.com/allenai/beaker-py/releases/tag/v1.22.0) - 2023-09-25
 
 ### Changed

diff --git a/beaker/services/experiment.py b/beaker/services/experiment.py
@@ -72,7 +72,38 @@ def _parse_create_args(
         spec: Optional[Union[ExperimentSpec, PathOrStr]] = kwargs.pop("spec", None)
         name: Optional[str] = kwargs.pop("name", None)
         workspace: Optional[Union[Workspace, str]] = kwargs.pop("workspace", None)
-        if len(args) == 2:
+        if len(args) == 3:
+            if name is not None or spec is not None or workspace is not None:
+                raise TypeError(
+                    "ExperimentClient.create() got an unexpected number of positional arguments"
+                )
+            if (
+                isinstance(args[0], str)
+                and isinstance(args[1], (ExperimentSpec, Path, str))
+                and isinstance(
+                    args[2],
+                    (
+                        str,
+                        Workspace,
+                    ),
+                )
+            ):
+                name, spec, workspace = args
+            elif (
+                isinstance(args[0], (ExperimentSpec, Path, str))
+                and isinstance(args[1], str)
+                and isinstance(
+                    args[2],
+                    (
+                        str,
+                        Workspace,
+                    ),
+                )
+            ):
+                spec, name, workspace = args
+            else:
+                raise TypeError("ExperimentClient.create() got an unexpected positional argument")
+        elif len(args) == 2:
             if name is not None or spec is not None:
                 raise TypeError(
                     "ExperimentClient.create() got an unexpected number of positional arguments"
@@ -90,14 +121,18 @@ def _parse_create_args(
                 name = args[0]
             else:
                 raise TypeError("ExperimentClient.create() got an unexpected positional argument")
-        else:
+        elif len(args) != 0:
             raise TypeError(
                 "ExperimentClient.create() got an unexpected number of positional arguments"
             )
+        elif spec is None:
+            raise TypeError("ExperimentClient.create() is missing 1 required argument: 'spec'")
+
         if kwargs:
             raise TypeError(
                 f"ExperimentClient.create() got unexpected keyword arguments {tuple(kwargs.keys())}"
             )
+
         assert spec is not None
         return spec, name, workspace
 

diff --git a/tests/experiment_test.py b/tests/experiment_test.py
@@ -39,6 +39,24 @@ def test_parse_create_args(client: Beaker):
     assert name == "my-experiment"
     assert spec is not None
 
+    spec, name, workspace = client.experiment._parse_create_args(
+        name="my-experiment",
+        spec=ExperimentSpec.new(docker_image="hello-world"),
+        workspace="ai2/petew",
+    )
+    assert workspace == "ai2/petew"
+    assert name == "my-experiment"
+    assert spec is not None
+
+    spec, name, workspace = client.experiment._parse_create_args(
+        "my-experiment",
+        ExperimentSpec.new(docker_image="hello-world"),
+        "ai2/petew",
+    )
+    assert workspace == "ai2/petew"
+    assert name == "my-experiment"
+    assert spec is not None
+
 
 def test_experiment_get(client: Beaker, hello_world_experiment_id: str):
     exp = client.experiment.get(hello_world_experiment_id)