chore: run formatter

codex-storage · Dec 20, 2024 · 99cc4d7 · 99cc4d7
1 parent 3c34929
commit 99cc4d7
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 35 deletions.
diff --git a/benchmarks/k8s/parameter_expander.py b/benchmarks/k8s/parameter_expander.py
@@ -23,8 +23,10 @@ def expand(parameters: Dict[str, Any], run_id: bool = False) -> List[Dict[str, A
         expanded_items = _expand_simple(simple)
     else:
         expanded_items = [
-            simple + constrained for simple, constrained in
-            itertools.product(_expand_simple(simple), _expand_constrained(constrained))
+            simple + constrained
+            for simple, constrained in itertools.product(
+                _expand_simple(simple), _expand_constrained(constrained)
+            )
         ]
 
     final_expansion = [dict(item, **fixed) for item in expanded_items]
@@ -36,27 +38,32 @@ def expand(parameters: Dict[str, Any], run_id: bool = False) -> List[Dict[str, A
     return final_expansion
 
 
-def _expand_simple(expandable: Dict[str, List[Any]]) -> List[List[Tuple[str, List[Any]]]]:
+def _expand_simple(
+    expandable: Dict[str, List[Any]],
+) -> List[List[Tuple[str, List[Any]]]]:
     keys = expandable.keys()
     return [
         list(zip(keys, list(value_set)))
         for value_set in itertools.product(*expandable.values())
     ]
 
 
-def _expand_constrained(constrained: Dict[str, List[Any]]) -> List[List[Tuple[str, List[Any]]]]:
+def _expand_constrained(
+    constrained: Dict[str, List[Any]],
+) -> List[List[Tuple[str, List[Any]]]]:
     return [
         expansion
         for k, v in constrained.items()
         for expansion in _expand_single_constraint(k, v)
     ]
 
 
-def _expand_single_constraint(combined_key: str,
-                              values: List[List[Any]]) -> List[List[Tuple[str, List[Any]]]]:
-    keys = combined_key[len('constrained__'):].split('_')
+def _expand_single_constraint(
+    combined_key: str, values: List[List[Any]]
+) -> List[List[Tuple[str, List[Any]]]]:
+    keys = combined_key[len("constrained__") :].split("_")
     if len(keys) < 2:
-        raise ValueError(f'Invalid combined key {combined_key}')
+        raise ValueError(f"Invalid combined key {combined_key}")
 
     normalized_values = [_normalize_values(value_set) for value_set in values]
 
@@ -68,10 +75,7 @@ def _expand_single_constraint(combined_key: str,
 
 
 def _normalize_values(values: List[Any | List[Any]]) -> List[List[Any]]:
-    return [
-        value if isinstance(value, list) else [value]
-        for value in values
-    ]
+    return [value if isinstance(value, list) else [value] for value in values]
 
 
 def normalize_argo_params(argo_params: List[Dict[str, Any]]) -> Dict[str, Any]:
@@ -85,7 +89,6 @@ def normalize_argo_params(argo_params: List[Dict[str, Any]]) -> Dict[str, Any]:
 
 
 if __name__ == "__main__":
-
     if len(sys.argv) < 2:
         print(f"Usage: {sys.argv[0]} '<json_string>'")
         sys.exit(1)

diff --git a/benchmarks/k8s/tests/test_parameter_expander.py b/benchmarks/k8s/tests/test_parameter_expander.py
@@ -5,12 +5,7 @@
 
 
 def test_should_expand_simple_parameter_lists():
-    matrix = {
-        "a": [1, 2],
-        "b": [3, 4],
-        "c": "foo",
-        "d": 5
-    }
+    matrix = {"a": [1, 2], "b": [3, 4], "c": "foo", "d": 5}
 
     assert expander.expand(matrix) == [
         {"a": 1, "b": 3, "c": "foo", "d": 5},
@@ -21,12 +16,7 @@ def test_should_expand_simple_parameter_lists():
 
 
 def test_should_add_run_id_when_requested():
-    matrix = {
-        "a": [1, 2],
-        "b": [3, 4],
-        "c": "foo",
-        "d": 5
-    }
+    matrix = {"a": [1, 2], "b": [3, 4], "c": "foo", "d": 5}
 
     assert expander.expand(matrix, run_id=True) == [
         {"a": 1, "b": 3, "c": "foo", "d": 5, "runId": 1},
@@ -37,13 +27,7 @@ def test_should_add_run_id_when_requested():
 
 
 def test_should_expand_constrained_parameter_pairs():
-    matrix = {
-        "constrained__att1_att2": [
-            [1, [2, 3]],
-            [[4, 5], 6]
-        ],
-        "b": [1, 2]
-    }
+    matrix = {"constrained__att1_att2": [[1, [2, 3]], [[4, 5], 6]], "b": [1, 2]}
 
     assert expander.expand(matrix) == [
         {"att1": 1, "att2": 2, "b": 1},
@@ -58,9 +42,11 @@ def test_should_expand_constrained_parameter_pairs():
 
 
 def test_should_normalize_simple_argo_parameter_list():
-    argo_params = json.loads('[{"name":"repetitions","value":"1"},{"name":"fileSize","value":"100MB"},'
-                             '{"name":"networkSize","value":"5"},{"name":"seeders","value":"1"},'
-                             '{"name":"seederSets","value":"1"},{"name":"maxExperimentDuration","value":"72h"}]')
+    argo_params = json.loads(
+        '[{"name":"repetitions","value":"1"},{"name":"fileSize","value":"100MB"},'
+        '{"name":"networkSize","value":"5"},{"name":"seeders","value":"1"},'
+        '{"name":"seederSets","value":"1"},{"name":"maxExperimentDuration","value":"72h"}]'
+    )
 
     assert normalize_argo_params(argo_params) == {
         "repetitions": 1,