fix partial wrappers and deferred annotations

python/cog/predictor.py

@@ -8,6 +8,7 @@
 import uuid
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
+from inspect import Signature
 from pathlib import Path
 from typing import (
     Any,
@@ -18,6 +19,7 @@
     Type,
     Union,
     cast,
+    get_type_hints,
 )
 
 try:
@@ -266,7 +268,7 @@ def validate_input_type(
     type: Type[Any],  # pylint: disable=redefined-builtin
     name: str,
 ) -> None:
-    if type is inspect.Signature.empty:
+    if type is 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."
         )
@@ -282,18 +284,20 @@ def validate_input_type(
             )
 
 
-def get_input_create_model_kwargs(signature: inspect.Signature) -> Dict[str, Any]:
+def get_input_create_model_kwargs(
+    signature: Signature, input_types: Dict[str, Any]
+) -> Dict[str, Any]:
     create_model_kwargs = {}
 
     order = 0
 
     for name, parameter in signature.parameters.items():
-        InputType = parameter.annotation
+        InputType = input_types.get(name, parameter.annotation)
 
         validate_input_type(InputType, name)
 
         # if no default is specified, create an empty, required input
-        if parameter.default is inspect.Signature.empty:
+        if parameter.default is Signature.empty:
             default = Input()
         else:
             default = parameter.default
@@ -354,25 +358,39 @@ class Input(BaseModel):
     predict = get_predict(predictor)
     signature = inspect.signature(predict)
 
+    try:
+        input_types = get_type_hints(predict)
+    except TypeError:
+        input_types = {}
+    if "return" in input_types:
+        del input_types["return"]
+
     return create_model(
         "Input",
         __config__=None,
         __base__=BaseInput,
         __module__=__name__,
         __validators__=None,
-        **get_input_create_model_kwargs(signature),
+        **get_input_create_model_kwargs(signature, input_types),
     )  # type: ignore
 
 
+def get_return_annotation(fn: Callable[..., Any]) -> "Type[Any]|Signature.empty":
+    try:
+        return get_type_hints(fn).get("return", Signature.empty)
+    except TypeError:
+        return inspect.signature(fn).return_annotation
+
+
 def get_output_type(predictor: BasePredictor) -> Type[BaseModel]:
     """
     Creates a Pydantic Output model from the return type annotation of a Predictor's predict() method.
     """
 
     predict = get_predict(predictor)
-    signature = inspect.signature(predict)
-    OutputType: Type[BaseModel]
-    if signature.return_annotation is inspect.Signature.empty:
+    maybe_output_type = get_return_annotation(predict)
+
+    if maybe_output_type is Signature.empty:
         raise TypeError(
             """You must set an output type. If your model can return multiple output types, you can explicitly set `Any` as the output type.
 
@@ -387,9 +405,8 @@ def predict(
         ...
 """
         )
-    else:
-        OutputType = signature.return_annotation
-
+    # we need the indirection to narrow the type to the one that's declared later
+    OutputType = maybe_output_type  # pylint: disable=invalid-name
     # The type that goes in the response is a list of the yielded type
     if get_origin(OutputType) is Iterator:
         # Annotated allows us to attach Field annotations to the list, which we use to mark that this is an iterator
@@ -452,13 +469,20 @@ class TrainingInput(BaseModel):
     train = get_train(predictor)
     signature = inspect.signature(train)
 
+    try:
+        input_types = get_type_hints(train)
+    except TypeError:
+        input_types = {}
+    if "return" in input_types:
+        del input_types["return"]
+
     return create_model(
         "TrainingInput",
         __config__=None,
         __base__=BaseInput,
         __module__=__name__,
         __validators__=None,
-        **get_input_create_model_kwargs(signature),
+        **get_input_create_model_kwargs(signature, input_types),
     )  # type: ignore
 
 
@@ -468,9 +492,9 @@ def get_training_output_type(predictor: BasePredictor) -> Type[BaseModel]:
     """
 
     train = get_train(predictor)
-    signature = inspect.signature(train)
 
-    if signature.return_annotation is inspect.Signature.empty:
+    TrainingOutputType = get_return_annotation(train)  # pylint: disable=invalid-name
+    if TrainingOutputType is Signature.empty:
         raise TypeError(
             """You must set an output type. If your model can return multiple output types, you can explicitly set `Any` as the output type.
 
@@ -485,8 +509,6 @@ def train(
         ...
 """
         )
-    else:
-        TrainingOutputType = signature.return_annotation
 
     name = (
         TrainingOutputType.__name__ if hasattr(TrainingOutputType, "__name__") else ""

test-integration/test_integration/fixtures/future-annotations-project/predict.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from cog import BasePredictor, Input
+
+
+class Predictor(BasePredictor):
+    def predict(self, input: str = Input(description="Who to greet")) -> str:
+        return "hello " + input

test-integration/test_integration/fixtures/partial-predict-project/cog.yaml

@@ -0,0 +1,4 @@
+build:
+  python_version: "3.8"
+predict: "predict.py:Predictor"
+train: "predict.py:train"

test-integration/test_integration/fixtures/partial-predict-project/predict.py

@@ -0,0 +1,39 @@
+import functools
+import inspect
+from typing import Any, Callable
+
+from cog import BasePredictor, Input
+
+
+class Predictor(BasePredictor):
+    def general(
+        self, prompt: str = Input(description="hi"), system_prompt: str = None
+    ) -> int:
+        return 1
+
+    def _remove(f: Callable, defaults: "dict[str, Any]") -> Callable:
+        # pylint: disable=no-self-argument
+        def wrapper(self, *args, **kwargs):
+            kwargs.update(defaults)
+            return f(self, *args, **kwargs)
+
+        # Update wrapper attributes for documentation, etc.
+        functools.update_wrapper(wrapper, f)
+
+        # for the purposes of inspect.signature as used by predictor.get_input_type,
+        # remove the argument (system_prompt)
+        sig = inspect.signature(f)
+        params = [p for name, p in sig.parameters.items() if name not in defaults]
+        wrapper.__signature__ = sig.replace(parameters=params)
+
+        # Return partialmethod, wrapper behaves correctly when part of a class
+        return functools.partialmethod(wrapper)
+
+    predict = _remove(general, {"system_prompt": ""})
+
+
+def _train(prompt: str = Input(description="hi"), system_prompt: str = None) -> int:
+    return 1
+
+
+train = functools.partial(_train, system_prompt="")

test-integration/test_integration/test_predict.py

@@ -288,3 +288,48 @@ def test_predict_path_list_input(tmpdir_factory):
     )
     assert "test1" in result.stdout
     assert "test2" in result.stdout
+
+
+def test_predict_works_with_deferred_annotations():
+    project_dir = Path(__file__).parent / "fixtures/future-annotations-project"
+
+    subprocess.check_call(
+        ["cog", "predict", "-i", "input=world"],
+        cwd=project_dir,
+        timeout=DEFAULT_TIMEOUT,
+    )
+
+
+def test_predict_works_with_partial_wrapper():
+    project_dir = Path(__file__).parent / "fixtures/partial-predict-project"
+
+    subprocess.check_call(
+        ["cog", "predict", "-i", "prompt=world"],
+        cwd=project_dir,
+        timeout=DEFAULT_TIMEOUT,
+    )
+    subprocess.check_call(
+        ["cog", "train", "-i", "prompt=world"],
+        cwd=project_dir,
+        timeout=DEFAULT_TIMEOUT,
+    )
+
+
+def test_predict_int_none_output():
+    project_dir = Path(__file__).parent / "fixtures/int-none-output-project"
+
+    subprocess.check_call(
+        ["cog", "predict"],
+        cwd=project_dir,
+        timeout=DEFAULT_TIMEOUT,
+    )
+
+
+def test_predict_string_none_output():
+    project_dir = Path(__file__).parent / "fixtures/string-none-output-project"
+
+    subprocess.check_call(
+        ["cog", "predict"],
+        cwd=project_dir,
+        timeout=DEFAULT_TIMEOUT,
+    )