pixeltable 0.2.24__py3-none-any.whl → 0.3.0__py3-none-any.whl

pixeltable/func/function.py

@@ -3,16 +3,22 @@ from __future__ import annotations
 import abc
 import importlib
 import inspect
-from typing import Any, Callable, Optional
+from copy import copy
+from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, cast
 
 import sqlalchemy as sql
+from typing_extensions import Self
 
 import pixeltable as pxt
+import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 
 from .globals import resolve_symbol
 from .signature import Signature
 
+if TYPE_CHECKING:
+    from .expr_template_function import ExprTemplate, ExprTemplateFunction
+
 
 class Function(abc.ABC):
     """Base class for Pixeltable's function interface.
@@ -22,30 +28,42 @@ class Function(abc.ABC):
     via the member self_path.
     """
 
-    signature: Signature
+    signatures: list[Signature]
     self_path: Optional[str]
     is_method: bool
     is_property: bool
     _conditional_return_type: Optional[Callable[..., ts.ColumnType]]
 
+    # We cache the overload resolutions in self._resolutions. This ensures that each resolution is represented
+    # globally by a single Python object. We do this dynamically rather than pre-constructing them in order to
+    # avoid circular complexity in the `Function` initialization logic.
+    __resolved_fns: list[Self]
+
     # Translates a call to this function with the given arguments to its SQLAlchemy equivalent.
     # Overriden for specific Function instances via the to_sql() decorator. The override must accept the same
     # parameter names as the original function. Each parameter is going to be of type sql.ColumnElement.
     _to_sql: Callable[..., Optional[sql.ColumnElement]]
 
-
     def __init__(
-        self, signature: Signature, self_path: Optional[str] = None, is_method: bool = False, is_property: bool = False
+        self,
+        signatures: list[Signature],
+        self_path: Optional[str] = None,
+        is_method: bool = False,
+        is_property: bool = False
     ):
         # Check that stored functions cannot be declared using `is_method` or `is_property`:
         assert not ((is_method or is_property) and self_path is None)
-        self.signature = signature
+        assert isinstance(signatures, list)
+        assert len(signatures) > 0
+        self.signatures = signatures
         self.self_path = self_path  # fully-qualified path to self
         self.is_method = is_method
         self.is_property = is_property
         self._conditional_return_type = None
         self._to_sql = self.__default_to_sql
 
+        self.__resolved_fns = []
+
     @property
     def name(self) -> str:
         assert self.self_path is not None
@@ -60,47 +78,198 @@ class Function(abc.ABC):
             return self.self_path[len(ptf_prefix):]
         return self.self_path
 
+    @property
+    def is_polymorphic(self) -> bool:
+        return len(self.signatures) > 1
+
+    @property
+    def signature(self) -> Signature:
+        assert not self.is_polymorphic
+        return self.signatures[0]
+
     @property
     def arity(self) -> int:
+        assert not self.is_polymorphic
         return len(self.signature.parameters)
 
+    def _docstring(self) -> Optional[str]:
+        return None
+
     def help_str(self) -> str:
-        return self.display_name + str(self.signature)
+        docstring = self._docstring()
+        display = self.display_name + str(self.signatures[0])
+        if docstring is None:
+            return display
+        return f'{display}\n\n{docstring}'
+
+    @property
+    def _resolved_fns(self) -> list[Self]:
+        """
+        Return the list of overload resolutions for this `Function`, constructing it first if necessary.
+        Each resolution is a new `Function` instance that retains just the single signature at index `signature_idx`,
+        and is otherwise identical to this `Function`.
+        """
+        if len(self.__resolved_fns) == 0:
+            # The list of overload resolutions hasn't been constructed yet; do so now.
+            if len(self.signatures) == 1:
+                # Only one signature: no need to construct separate resolutions
+                self.__resolved_fns.append(self)
+            else:
+                # Multiple signatures: construct a resolution for each signature
+                for idx in range(len(self.signatures)):
+                    resolution = cast(Self, copy(self))
+                    resolution.signatures = [self.signatures[idx]]
+                    resolution.__resolved_fns = [resolution]  # Resolves to itself
+                    resolution._update_as_overload_resolution(idx)
+                    self.__resolved_fns.append(resolution)
+
+        return self.__resolved_fns
+
+    @property
+    def _has_resolved_fns(self) -> bool:
+        """
+        Returns true if the resolved_fns for this `Function` have been constructed (i.e., if self._resolved_fns
+        has been accessed).
+        """
+        return len(self.__resolved_fns) > 0
+
+    def _update_as_overload_resolution(self, signature_idx: int) -> None:
+        """
+        Subclasses must implement this in order to do any additional work when creating a resolution, beyond
+        simply updating `self.signatures`.
+        """
+        raise NotImplementedError()
 
     def __call__(self, *args: Any, **kwargs: Any) -> 'pxt.exprs.FunctionCall':
         from pixeltable import exprs
-        bound_args = self.signature.py_signature.bind(*args, **kwargs)
-        self.validate_call(bound_args.arguments)
-        return exprs.FunctionCall(self, bound_args.arguments)
+
+        resolved_fn, bound_args = self._bind_to_matching_signature(args, kwargs)
+        return_type = resolved_fn.call_return_type(args, kwargs)
+        return exprs.FunctionCall(resolved_fn, bound_args, return_type)
+
+    def _bind_to_matching_signature(self, args: Sequence[Any], kwargs: dict[str, Any]) -> tuple[Self, dict[str, Any]]:
+        result: int = -1
+        bound_args: Optional[dict[str, Any]] = None
+        assert len(self.signatures) > 0
+        if len(self.signatures) == 1:
+            # Only one signature: call _bind_to_signature() and surface any errors directly
+            result = 0
+            bound_args = self._bind_to_signature(0, args, kwargs)
+        else:
+            # Multiple signatures: try each signature in declaration order and trap any errors.
+            # If none of them succeed, raise a generic error message.
+            for i in range(len(self.signatures)):
+                try:
+                    bound_args = self._bind_to_signature(i, args, kwargs)
+                except (TypeError, excs.Error):
+                    continue
+                result = i
+                break
+            if result == -1:
+                raise excs.Error(f'Function {self.name!r} has no matching signature for arguments')
+        assert result >= 0
+        assert bound_args is not None
+        return self._resolved_fns[result], bound_args
+
+    def _bind_to_signature(self, signature_idx: int, args: Sequence[Any], kwargs: dict[str, Any]) -> dict[str, Any]:
+        from pixeltable import exprs
+
+        signature = self.signatures[signature_idx]
+        bound_args = signature.py_signature.bind(*args, **kwargs).arguments
+        self._resolved_fns[signature_idx].validate_call(bound_args)
+        exprs.FunctionCall.normalize_args(self.name, signature, bound_args)
+        return bound_args
 
     def validate_call(self, bound_args: dict[str, Any]) -> None:
         """Override this to do custom validation of the arguments"""
-        pass
+        assert not self.is_polymorphic
 
-    def call_return_type(self, kwargs: dict[str, Any]) -> ts.ColumnType:
+    def call_return_type(self, args: Sequence[Any], kwargs: dict[str, Any]) -> ts.ColumnType:
         """Return the type of the value returned by calling this function with the given arguments"""
+        assert not self.is_polymorphic
         if self._conditional_return_type is None:
             return self.signature.return_type
-        bound_args = self.signature.py_signature.bind(**kwargs)
+        bound_args = self.signature.py_signature.bind(*args, **kwargs).arguments
         kw_args: dict[str, Any] = {}
         sig = inspect.signature(self._conditional_return_type)
         for param in sig.parameters.values():
-            if param.name in bound_args.arguments:
-                kw_args[param.name] = bound_args.arguments[param.name]
+            if param.name in bound_args:
+                kw_args[param.name] = bound_args[param.name]
         return self._conditional_return_type(**kw_args)
 
     def conditional_return_type(self, fn: Callable[..., ts.ColumnType]) -> Callable[..., ts.ColumnType]:
         """Instance decorator for specifying a conditional return type for this function"""
         # verify that call_return_type only has parameters that are also present in the signature
-        sig = inspect.signature(fn)
-        for param in sig.parameters.values():
-            if param.name not in self.signature.parameters:
-                raise ValueError(f'`conditional_return_type` has parameter `{param.name}` that is not in the signature')
+        fn_sig = inspect.signature(fn)
+        for param in fn_sig.parameters.values():
+            for self_sig in self.signatures:
+                if param.name not in self_sig.parameters:
+                    raise ValueError(f'`conditional_return_type` has parameter `{param.name}` that is not in a signature')
         self._conditional_return_type = fn
         return fn
 
+    def using(self, **kwargs: Any) -> 'ExprTemplateFunction':
+        from .expr_template_function import ExprTemplateFunction
+
+        assert len(self.signatures) > 0
+        if len(self.signatures) == 1:
+            # Only one signature: call _bind_and_create_template() and surface any errors directly
+            template = self._bind_and_create_template(kwargs)
+            return ExprTemplateFunction([template])
+        else:
+            # Multiple signatures: iterate over each signature and generate a template for each
+            # successful binding. If there are no successful bindings, raise a generic error.
+            # (Note that the resulting ExprTemplateFunction may have strictly fewer signatures than
+            # this Function, in the event that only some of the signatures are successfully bound.)
+            templates: list['ExprTemplate'] = []
+            for i in range(len(self.signatures)):
+                try:
+                    template = self._resolved_fns[i]._bind_and_create_template(kwargs)
+                    templates.append(template)
+                except (TypeError, excs.Error):
+                    continue
+            if len(templates) == 0:
+                raise excs.Error(f'Function {self.name!r} has no matching signature for arguments')
+            return ExprTemplateFunction(templates)
+
+    def _bind_and_create_template(self, kwargs: dict[str, Any]) -> 'ExprTemplate':
+        from pixeltable import exprs
+
+        from .expr_template_function import ExprTemplate
+
+        assert not self.is_polymorphic
+
+        # Resolve each kwarg into a parameter binding
+        bindings: dict[str, exprs.Expr] = {}
+        for k, v in kwargs.items():
+            if k not in self.signature.parameters:
+                raise excs.Error(f'Unknown parameter: {k}')
+            param = self.signature.parameters[k]
+            expr = exprs.Expr.from_object(v)
+            if not param.col_type.is_supertype_of(expr.col_type):
+                raise excs.Error(f'Expected type `{param.col_type}` for parameter `{k}`; got `{expr.col_type}`')
+            bindings[k] = v  # Use the original value, not the Expr (The Expr is only for validation)
+
+        residual_params = [
+            p for p in self.signature.parameters.values() if p.name not in bindings
+        ]
+
+        # Bind each remaining parameter to a like-named variable
+        for param in residual_params:
+            bindings[param.name] = exprs.Variable(param.name, param.col_type)
+
+        return_type = self.call_return_type([], bindings)
+        call = exprs.FunctionCall(self, bindings, return_type)
+
+        # Construct the (n-k)-ary signature of the new function. We use `call.col_type` for this, rather than
+        # `self.signature.return_type`, because the return type of the new function may be specialized via a
+        # conditional return type.
+        new_signature = Signature(call.col_type, residual_params, self.signature.is_batched)
+
+        return ExprTemplate(call, new_signature)
+
     @abc.abstractmethod
-    def exec(self, *args: Any, **kwargs: Any) -> Any:
+    def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
         """Execute the function with the given arguments and return the result."""
         pass
 
@@ -122,19 +291,24 @@ class Function(abc.ABC):
         """Print source code"""
         print('source not available')
 
-    def as_dict(self) -> dict:
+    def as_dict(self) -> dict[str, Any]:
         """
         Return a serialized reference to the instance that can be passed to json.dumps() and converted back
         to an instance with from_dict().
         Subclasses can override _as_dict().
         """
+        # We currently only ever serialize a function that has a specific signature (not a polymorphic form).
+        assert not self.is_polymorphic
         classpath = f'{self.__class__.__module__}.{self.__class__.__qualname__}'
         return {'_classpath': classpath, **self._as_dict()}
 
     def _as_dict(self) -> dict:
-        """Default serialization: store the path to self (which includes the module path)"""
+        """Default serialization: store the path to self (which includes the module path) and signature."""
         assert self.self_path is not None
-        return {'path': self.self_path}
+        return {
+            'path': self.self_path,
+            'signature': self.signature.as_dict(),
+        }
 
     @classmethod
     def from_dict(cls, d: dict) -> Function:
@@ -145,15 +319,45 @@ class Function(abc.ABC):
         module_path, class_name = d['_classpath'].rsplit('.', 1)
         class_module = importlib.import_module(module_path)
         func_class = getattr(class_module, class_name)
+        assert isinstance(func_class, type) and issubclass(func_class, Function)
         return func_class._from_dict(d)
 
     @classmethod
     def _from_dict(cls, d: dict) -> Function:
         """Default deserialization: load the symbol indicated by the stored symbol_path"""
         assert 'path' in d and d['path'] is not None
+        assert 'signature' in d and d['signature'] is not None
         instance = resolve_symbol(d['path'])
         assert isinstance(instance, Function)
-        return instance
+
+        # Load the signature from the DB and check that it is still valid (i.e., is still consistent with a signature
+        # in the code).
+        signature = Signature.from_dict(d['signature'])
+        idx = instance.__find_matching_overload(signature)
+        if idx is None:
+            # No match; generate an informative error message.
+            signature_note_str = 'any of its signatures' if instance.is_polymorphic else 'its signature as'
+            instance_signature_str = (
+                f'{len(instance.signatures)} signatures' if instance.is_polymorphic else str(instance.signature)
+            )
+            # TODO: Handle this more gracefully (instead of failing the DB load, allow the DB load to succeed, but
+            #       mark any enclosing FunctionCall as unusable). It's the same issue as dealing with a renamed UDF or
+            #       FunctionCall return type mismatch.
+            raise excs.Error(
+                f'The signature stored in the database for the UDF `{instance.self_path}` no longer matches '
+                f'{signature_note_str} as currently defined in the code.\nThis probably means that the code for '
+                f'`{instance.self_path}` has changed in a backward-incompatible way.\n'
+                f'Signature in database: {signature}\n'
+                f'Signature in code: {instance_signature_str}'
+            )
+        # We found a match; specialize to the appropriate overload resolution (non-polymorphic form) and return that.
+        return instance._resolved_fns[idx]
+
+    def __find_matching_overload(self, sig: Signature) -> Optional[int]:
+        for idx, overload_sig in enumerate(self.signatures):
+            if sig.is_consistent_with(overload_sig):
+                return idx
+        return None
 
     def to_store(self) -> tuple[dict, bytes]:
         """

pixeltable/func/function_registry.py

@@ -13,6 +13,7 @@ import pixeltable.env as env
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 from pixeltable.metadata import schema
+
 from .function import Function
 
 _logger = logging.getLogger('pixeltable')
@@ -68,7 +69,7 @@ class FunctionRegistry:
             raise excs.Error(f'A UDF with that name already exists: {fqn}')
         self.module_fns[fqn] = fn
         if fn.is_method or fn.is_property:
-            base_type = fn.signature.parameters_by_pos[0].col_type.type_enum
+            base_type = fn.signatures[0].parameters_by_pos[0].col_type.type_enum
             if base_type not in self.type_methods:
                 self.type_methods[base_type] = {}
             if fn.name in self.type_methods[base_type]:

pixeltable/func/query_template_function.py

@@ -1,23 +1,27 @@
 from __future__ import annotations
 
 import inspect
-from typing import Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, overload
 
 import sqlalchemy as sql
 
-import pixeltable as pxt
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
 from pixeltable import exprs
 
 from .function import Function
 from .signature import Signature
 
+if TYPE_CHECKING:
+    from pixeltable import DataFrame
+
 
 class QueryTemplateFunction(Function):
     """A parameterized query/DataFrame from which an executable DataFrame is created with a function call."""
 
     @classmethod
     def create(
-        cls, template_callable: Callable, param_types: Optional[list[pxt.ColumnType]], path: str, name: str
+        cls, template_callable: Callable, param_types: Optional[list[ts.ColumnType]], path: str, name: str
     ) -> QueryTemplateFunction:
         # we need to construct a template df and a signature
         py_sig = inspect.signature(template_callable)
@@ -29,14 +33,15 @@ class QueryTemplateFunction(Function):
         from pixeltable import DataFrame
         assert isinstance(template_df, DataFrame)
         # we take params and return json
-        sig = Signature(return_type=pxt.JsonType(), parameters=params)
+        sig = Signature(return_type=ts.JsonType(), parameters=params)
         return QueryTemplateFunction(template_df, sig, path=path, name=name)
 
     def __init__(
-            self, template_df: Optional['pxt.DataFrame'], sig: Optional[Signature], path: Optional[str] = None,
+            self, template_df: Optional['DataFrame'], sig: Signature, path: Optional[str] = None,
             name: Optional[str] = None,
     ):
-        super().__init__(sig, self_path=path)
+        assert sig is not None
+        super().__init__([sig], self_path=path)
         self.self_name = name
         self.template_df = template_df
 
@@ -48,16 +53,20 @@ class QueryTemplateFunction(Function):
         # convert defaults to Literals
         self.defaults: dict[str, exprs.Literal] = {}  # key: param name, value: default value converted to a Literal
         param_types = self.template_df.parameters()
-        for param in [p for p in self.signature.parameters.values() if p.has_default()]:
+        for param in [p for p in sig.parameters.values() if p.has_default()]:
             assert param.name in param_types
             param_type = param_types[param.name]
             literal_default = exprs.Literal(param.default, col_type=param_type)
             self.defaults[param.name] = literal_default
 
+    def _update_as_overload_resolution(self, signature_idx: int) -> None:
+        pass  # only one signature supported for QueryTemplateFunction
+
     def set_conn(self, conn: Optional[sql.engine.Connection]) -> None:
         self.conn = conn
 
-    def exec(self, *args: Any, **kwargs: Any) -> Any:
+    def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
+        assert not self.is_polymorphic
         bound_args = self.signature.py_signature.bind(*args, **kwargs).arguments
         # apply defaults, otherwise we might have Parameters left over
         bound_args.update(
@@ -75,9 +84,42 @@ class QueryTemplateFunction(Function):
         return self.self_name
 
     def _as_dict(self) -> dict:
-        return {'name': self.name, 'signature': self.signature.as_dict(), 'df': self.template_df.as_dict()}
+        return {'name': self.name, 'signature': self.signatures[0].as_dict(), 'df': self.template_df.as_dict()}
 
     @classmethod
     def _from_dict(cls, d: dict) -> Function:
         from pixeltable.dataframe import DataFrame
         return cls(DataFrame.from_dict(d['df']), Signature.from_dict(d['signature']), name=d['name'])
+
+
+@overload
+def query(self, py_fn: Callable) -> QueryTemplateFunction: ...
+
+@overload
+def query(
+    self, *, param_types: Optional[list[ts.ColumnType]] = None
+) -> Callable[[Callable], QueryTemplateFunction]: ...
+
+def query(*args: Any, **kwargs: Any) -> Any:
+    def make_query_template(
+        py_fn: Callable, param_types: Optional[list[ts.ColumnType]]
+    ) -> QueryTemplateFunction:
+        if py_fn.__module__ != '__main__' and py_fn.__name__.isidentifier():
+            # this is a named function in a module
+            function_path = f'{py_fn.__module__}.{py_fn.__qualname__}'
+        else:
+            function_path = None
+        query_name = py_fn.__name__
+        query_fn = QueryTemplateFunction.create(
+            py_fn, param_types=param_types, path=function_path, name=query_name)
+        return query_fn
+
+        # TODO: verify that the inferred return type matches that of the template
+        # TODO: verify that the signature doesn't contain batched parameters
+
+    if len(args) == 1:
+        assert len(kwargs) == 0 and callable(args[0])
+        return make_query_template(args[0], None)
+    else:
+        assert len(args) == 0 and len(kwargs) == 1 and 'param_types' in kwargs
+        return lambda py_fn: make_query_template(py_fn, kwargs['param_types'])

pixeltable/func/signature.py

@@ -91,6 +91,7 @@ class Signature:
         self.parameters_by_pos = parameters.copy()
         self.constant_parameters = [p for p in parameters if not p.is_batched]
         self.batched_parameters = [p for p in parameters if p.is_batched]
+        self.required_parameters = [p for p in parameters if not p.has_default()]
         self.py_signature = inspect.Signature([p.to_py_param() for p in self.parameters_by_pos])
 
     def get_return_type(self) -> ts.ColumnType:
@@ -110,6 +111,38 @@ class Signature:
         parameters = [Parameter.from_dict(param_dict) for param_dict in d['parameters']]
         return cls(ts.ColumnType.from_dict(d['return_type']), parameters, d['is_batched'])
 
+    def is_consistent_with(self, other: Signature) -> bool:
+        """
+        Returns True if this signature is consistent with the other signature.
+        S is consistent with T if we could safely replace S by T in any call where S is used. Specifically:
+        (i) S.return_type is a supertype of T.return_type
+        (ii) For each parameter p in S, there is a parameter q in T such that:
+            - p and q have the same name and kind
+            - q.col_type is a supertype of p.col_type
+        (iii) For each *required* parameter q in T, there is a parameter p in S with the same name (in which
+            case the kinds and types must also match, by condition (ii)).
+        """
+        # Check (i)
+        if not self.get_return_type().is_supertype_of(other.get_return_type(), ignore_nullable=True):
+            return False
+
+        # Check (ii)
+        for param_name, param in self.parameters.items():
+            if param_name not in other.parameters:
+                return False
+            other_param = other.parameters[param_name]
+            if (param.kind != other_param.kind or
+                (param.col_type is None) != (other_param.col_type is None) or  # this can happen if they are varargs
+                param.col_type is not None and not other_param.col_type.is_supertype_of(param.col_type, ignore_nullable=True)):
+                return False
+
+        # Check (iii)
+        for other_param in other.required_parameters:
+            if other_param.name not in self.parameters:
+                return False
+
+        return True
+
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, Signature):
             return False
@@ -155,8 +188,12 @@ class Signature:
 
     @classmethod
     def create_parameters(
-            cls, py_fn: Optional[Callable] = None, py_params: Optional[list[inspect.Parameter]] = None,
-            param_types: Optional[list[ts.ColumnType]] = None
+        cls,
+        py_fn: Optional[Callable] = None,
+        py_params: Optional[list[inspect.Parameter]] = None,
+        param_types: Optional[list[ts.ColumnType]] = None,
+        type_substitutions: Optional[dict] = None,
+        is_cls_method: bool = False
     ) -> list[Parameter]:
         assert (py_fn is None) != (py_params is None)
         if py_fn is not None:
@@ -164,7 +201,12 @@ class Signature:
             py_params = list(sig.parameters.values())
         parameters: list[Parameter] = []
 
+        if type_substitutions is None:
+            type_substitutions = {}
+
         for idx, param in enumerate(py_params):
+            if is_cls_method and idx == 0:
+                continue  # skip 'self' or 'cls' parameter
             if param.name in cls.SPECIAL_PARAM_NAMES:
                 raise excs.Error(f"'{param.name}' is a reserved parameter name")
             if param.kind == inspect.Parameter.VAR_POSITIONAL or param.kind == inspect.Parameter.VAR_KEYWORD:
@@ -178,7 +220,12 @@ class Signature:
                 param_type = param_types[idx]
                 is_batched = False
             else:
-                param_type, is_batched = cls._infer_type(param.annotation)
+                py_type: Optional[type]
+                if param.annotation in type_substitutions:
+                    py_type = type_substitutions[param.annotation]
+                else:
+                    py_type = param.annotation
+                param_type, is_batched = cls._infer_type(py_type)
                 if param_type is None:
                     raise excs.Error(f'Cannot infer pixeltable type for parameter {param.name}')
 
@@ -189,18 +236,29 @@ class Signature:
 
     @classmethod
     def create(
-        cls, py_fn: Callable,
+        cls,
+        py_fn: Callable,
         param_types: Optional[list[ts.ColumnType]] = None,
-        return_type: Optional[ts.ColumnType] = None
+        return_type: Optional[ts.ColumnType] = None,
+        type_substitutions: Optional[dict] = None,
+        is_cls_method: bool = False
     ) -> Signature:
         """Create a signature for the given Callable.
         Infer the parameter and return types, if none are specified.
         Raises an exception if the types cannot be inferred.
         """
-        parameters = cls.create_parameters(py_fn=py_fn, param_types=param_types)
+        if type_substitutions is None:
+            type_substitutions = {}
+
+        parameters = cls.create_parameters(py_fn=py_fn, param_types=param_types, is_cls_method=is_cls_method, type_substitutions=type_substitutions)
         sig = inspect.signature(py_fn)
         if return_type is None:
-            return_type, return_is_batched = cls._infer_type(sig.return_annotation)
+            py_type: Optional[type]
+            if sig.return_annotation in type_substitutions:
+                py_type = type_substitutions[sig.return_annotation]
+            else:
+                py_type = sig.return_annotation
+            return_type, return_is_batched = cls._infer_type(py_type)
             if return_type is None:
                 raise excs.Error('Cannot infer pixeltable return type')
         else: