improve type hints

svpino · Dec 7, 2024 · 49dbbef · 49dbbef
1 parent 5c8ac36
commit 49dbbef
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/pipelines/inference.py b/pipelines/inference.py
@@ -88,7 +88,7 @@ def load_context(self, context: PythonModelContext) -> None:
     def predict(
         self,
         context: PythonModelContext,  # noqa: ARG002
-        model_input: pd.DataFrame | list | dict,
+        model_input: pd.DataFrame | list[dict[str, Any]] | dict[str, Any] | list[Any],
         params: dict[str, Any] | None = None,
     ) -> list:
         """Handle the request received from the client.

diff --git a/tests/test_inference.py b/tests/test_inference.py
@@ -140,9 +140,9 @@ def test_predict_return_empty_list_on_invalid_prediction(model, monkeypatch):
     "input_data",
     [
         pd.DataFrame([{"island": "Torgersen", "culmen_length_mm": 39.1}]),
-        ["Torgersen", 39.1],
-        {"island": "Torgersen", "culmen_length_mm": 39.1},
         [{"island": "Torgersen", "culmen_length_mm": 39.1}],
+        {"island": "Torgersen", "culmen_length_mm": 39.1},
+        ["Torgersen", 39.1],
     ],
 )
 def test_predict(model, monkeypatch, input_data):