Infer return type annotation from other function's annotation

Question

I have a function with a complex return type annotation:

from typing import (Union, List)

# The -> Union[…] is actually longer than this example:
def parse(what: str) -> Union[None, int, bool, float, complex, List[int]]:
  # do stuff and return the object as needed

def parse_from_something(which: SomeType) -> ????:
  return parse(which.extract_string())

# …

class SomeType:
  def extract_string(self) -> str:
     # do stuff and return str

How can I type-annotate parse_from_something so that it is annotated to return the same types as parse, without repeating them?

The problem I’m solving here is that one function is subject to change, but there’s wrappers around it that will always return the identical set of types. I don’t want to duplicate code, and because this is a refactoring and after-the-fact type annotation effort, I need to assume I’ll remove possible return types from parse in the future, and a static type checker might not realize that parse_from_something can no longer return these.

Asked By: Marcus Müller

||

Source

Answer 1

You could use a type alias, e.g.:

from typing import (Any, Union, List)


ParseResultType = Union[None, int, bool, float, complex, List[int]]


def parse(what: str) -> ParseResultType:
    # do stuff and return the object as needed


def parse_from_something(which: Any) -> ParseResultType:
    return parse(which.extract_string())


class SomeType:
    def extract_string(self) -> str:
        # do stuff and return str

Answered By: erny

Answer 2

What you describe is not possible by design.

From PEP 484:

It is recommended but not required that checked functions have annotations for all arguments and the return type. For a checked function, the default annotation for arguments and for the return type is Any. An exception is the first argument of instance and class methods. […]

So, if you omit the return type, it’s Any for external observer. No othe treatment is possible.

You have now two options. The first one (preferred, because it gives you type safety and gives your annotations better semantic meaning) is using an alias and just putting it everywhere when you want to say "same type as in X". It obviously scales not only to function return types, but also to class attributes, global constants, function and method arguments, etc. This approach is explained well in @erny answer.

Another option is using some trick similar to this answer (but please rename it to snake_case for PEP8 compliance – it’s still a function). It introduces a helper decorator to lie a little to your typechecker. When you apply this withParameterAndReturnTypesOf decorator, you basically say: "ignore what happens in function body, treat my func as a twin of the referenced one". It is a beautiful cheat, but this way you don’t get the real safety: your implementation may start doing more work and deviate from the contract you signed with the typechecker. E.g. the following will pass:

def parse(what: str) -> Union[None, int, bool, float, complex, List[int]]:
  # do stuff and return the object as needed

@dangerous_with_parameters_and_return_type_of(parse)
def parse_from_something(which: SomeType):
    if which is None: return NotImplemented  # not covered by the union
    return parse(which.extract_string())

If you were using the aliased union, you’d get an error from mypy that return type is incompatible.

Answered By: STerliakov

Answer 3

This is possible using some type variables and a decorator that inserts a particular return type into the parameter specification of your function.

Start with the following setup:

from typing import ParamSpec, TypeVar, Union, List, Any, Callable

P = ParamSpec("P")
R = TypeVar("R")

def parse(
    what: str,
) -> Union[None, int, bool, float, complex, List[int]]:
    ...

def identity(x: R) -> R:
    return x

def returns_like(func: Callable[..., R]) -> Callable[[Callable[P, Any]], Callable[P, R]]:
    return identity

You can then define parse_from_something as

@returns_like(parse)
def parse_from_something(which: Any) -> Any:
    return parse(which.extract_string())

reveal_type(parse_from_something)
reveal_type(parse_from_something(123))

This passes type checking under mypy --strict and pyright, and outputs the following revealed types:

# mypy --strict
note: Revealed type is "def (which: Any) -> Union[None, builtins.int, builtins.float, builtins.complex, builtins.list[builtins.int]]"
note: Revealed type is "Union[None, builtins.int, builtins.float, builtins.complex, builtins.list[builtins.int]]"
Success: no issues found in 1 source file

# pyright
information: Type of "parse_from_something" is "(which: Any) -> (int | bool | float | complex | List[int] | None)"
information: Type of "parse_from_something(123)" is "int | bool | float | complex | List[int] | None"
0 errors, 0 warnings, 2 informations

Try it yourself online: mypy playground and pyright playground.

Do you that inside of parse_from_something, your type checker won’t be able to validate that whatever you return is compatible with the return type of parse. For example, if you added a return {}, your type checker would remain silent instead of flagging it as a type error. You assume responsibility for making sure that parse_from_something is well behaved.

P.S. If you’re using an old version of Python that does not have typing.TypeVar or typing.ParamSpec, you can use typing-extensions to backport them. Alternatively, if you’re using Python 3.12+, you can replace these with PEP 695 Type Parameter Syntax.

Answered By: Brian61354270

Infer return type annotation from other function's annotation

Question:

Answers: