Allow Union and List input types (#1311)

* Allow Union and List input types

Signed-off-by: Mattt Zmuda <mattt@replicate.com>

* Update TypeError messages to note support for Union and List types

Signed-off-by: Mattt Zmuda <mattt@replicate.com>

---------

Signed-off-by: Mattt Zmuda <mattt@replicate.com>

Author: mattt

python/cog/predictor.py

@@ -214,6 +214,19 @@ def get_predict(predictor: Any) -> Callable:
         return predictor.predict
     return predictor
 
+def validate_input_type(type: Type, name: str) -> None:
+    if type is inspect.Signature.empty:
+            raise TypeError(
+                f"No input type provided for parameter `{name}`. Supported input types are: {readable_types_list(ALLOWED_INPUT_TYPES)}, or a Union or List of those types."
+            )
+    elif type not in ALLOWED_INPUT_TYPES:
+        if hasattr(type, "__origin__") and (type.__origin__ is Union or type.__origin__ is list):
+            for t in get_args(type):
+                validate_input_type(t, name)
+        else:
+            raise TypeError(
+                f"Unsupported input type {human_readable_type_name(type)} for parameter `{name}`. Supported input types are: {readable_types_list(ALLOWED_INPUT_TYPES)}, or a Union or List of those types."
+            )
 
 def get_input_type(predictor: BasePredictor) -> Type[BaseInput]:
     """
@@ -238,14 +251,7 @@ class Input(BaseModel):
     for name, parameter in signature.parameters.items():
         InputType = parameter.annotation
 
-        if InputType is inspect.Signature.empty:
-            raise TypeError(
-                f"No input type provided for parameter `{name}`. Supported input types are: {readable_types_list(ALLOWED_INPUT_TYPES)}."
-            )
-        elif InputType not in ALLOWED_INPUT_TYPES:
-            raise TypeError(
-                f"Unsupported input type {human_readable_type_name(InputType)} for parameter `{name}`. Supported input types are: {readable_types_list(ALLOWED_INPUT_TYPES)}."
-            )
+        validate_input_type(InputType, name)
 
         # if no default is specified, create an empty, required input
         if parameter.default is inspect.Signature.empty:

python/tests/server/fixtures/input_union_integer_or_list_of_integers.py

@@ -0,0 +1,10 @@
+from cog import BasePredictor, Path
+
+from typing import List, Union
+
+class Predictor(BasePredictor):
+    def predict(self, args: Union[int, List[int]]) -> int:
+        if isinstance(args, int):
+            return args
+        else:
+            return sum(args)

python/tests/server/fixtures/input_union_string_or_list_of_strings.py

@@ -0,0 +1,10 @@
+from cog import BasePredictor, Path
+
+from typing import List, Union
+
+class Predictor(BasePredictor):
+    def predict(self, args: Union[str, List[str]]) -> str:
+        if isinstance(args, str):
+            return args
+        else:
+            return "".join(args)

python/tests/server/test_http_input.py

@@ -30,13 +30,6 @@ def test_empty_input(client, match):
     assert resp.json() == match({"status": "succeeded", "output": "foobar"})
 
 
-@uses_predictor("input_string")
-def test_good_str_input(client, match):
-    resp = client.post("/predictions", json={"input": {"text": "baz"}})
-    assert resp.status_code == 200
-    assert resp.json() == match({"status": "succeeded", "output": "baz"})
-
-
 @uses_predictor("input_integer")
 def test_good_int_input(client, match):
     resp = client.post("/predictions", json={"input": {"num": 3}})
@@ -212,6 +205,39 @@ def test_choices_int(client):
     assert resp.status_code == 422
 
 
+@uses_predictor("input_union_string_or_list_of_strings")
+def test_union_strings(client):
+    resp = client.post("/predictions", json={"input": {"args": "abc"}})
+    assert resp.status_code == 200
+    assert resp.json()["output"] == "abc"
+
+    resp = client.post("/predictions", json={"input": {"args": ["a", "b", "c"]}})
+    assert resp.status_code == 200
+    assert resp.json()["output"] == "abc"
+
+    # FIXME: Numbers are successfully cast to strings, but maybe shouldn't be
+    # resp = client.post("/predictions", json={"input": {"args": 123}})
+    # assert resp.status_code == 422
+    # resp = client.post("/predictions", json={"input": {"args": [1, 2, 3]}})
+    # assert resp.status_code == 422
+
+
+@uses_predictor("input_union_integer_or_list_of_integers")
+def test_union_integers(client):
+    resp = client.post("/predictions", json={"input": {"args": 123}})
+    assert resp.status_code == 200
+    assert resp.json()["output"] == 123
+
+    resp = client.post("/predictions", json={"input": {"args": [1, 2, 3]}})
+    assert resp.status_code == 200
+    assert resp.json()["output"] == 6
+
+    resp = client.post("/predictions", json={"input": {"args": "abc"}})
+    assert resp.status_code == 422
+    resp = client.post("/predictions", json={"input": {"args": ["a", "b", "c"]}})
+    assert resp.status_code == 422
+
+
 def test_untyped_inputs():
     with pytest.raises(TypeError):
         make_client("input_untyped")