TaskSpec.with_contraint takes kwargs

CHANGELOG.md

@@ -9,6 +9,10 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Changed
+
+- `TaskSpec.with_contraint` now takes `**kwargs` instead of two positional arguments.
+
 ## [v1.11.3](https://github.com/allenai/beaker-py/releases/tag/v1.11.3) - 2022-10-28
 
 ### Added

beaker/data_model/experiment_spec.py

@@ -583,25 +583,24 @@ def with_env_var(
             },
         )
 
-    def with_constraint(self, constraint_type: str, constraint_values: List[str]) -> "TaskSpec":
+    def with_constraint(self, **kwargs: List[str]) -> "TaskSpec":
         """
         Return a new :class:`TaskSpec` with the given :data:`constraint`.
 
-        :param constraint_type: The type of constraint, e.g.  :data:`cluster`.
-        :param constraint_values: The list of constraint values, such as a list of cluster names.
+        :param kwargs: Constraint name, constraint values.
 
         :examples:
 
         >>> task_spec = TaskSpec.new(
         ...     "hello-world",
         ...     docker_image="hello-world",
-        ... ).with_constraint('cluster', ['ai2/cpu-cluster'])
+        ... ).with_constraint(cluster=['ai2/cpu-cluster'])
         >>> assert task_spec.constraints['cluster'] == ['ai2/cpu-cluster']
         """
         return self.copy(
             deep=True,
             update={
-                "constraints": {**(self.constraints or {}), constraint_type: constraint_values},
+                "constraints": {**(self.constraints or {}), **kwargs},
             },
         )