How to type a Python function the same way as another function?

448 views Asked by At

For writing a wrapper around an existing function, I want to that wrapper to have the same, or very similar, type.

For example:

import os

def my_open(*args, **kwargs):
    return os.open(*args, **kwargs)

Tye type signature for os.open() is complex and may change over time as its functionality and typings evolve, so I do not want to copy-paste the type signature of os.open() into my code. Instead, I want to infer the type for my_open(), so that it "copies" the type of os.open()'s parameters and return values.

my_open() shall have the same type as the wrapped function os.open().


I would like to do the same thing with a decorated function:

@contextmanager
def scoped_open(*args, **kwargs):
  """Like `os.open`, but as a `contextmanager` yielding the FD.
  """
  fd = os.open(*args, **kwargs)
  try:
    yield fd
  finally:
    os.close(fd)

Here, the inferred function arguments of scoped_open() shall be the same ones as os.open(), but the return type shall be a Generator of the inferred return type of os.open() (currently int, but again I do not wish to copy-paste that int).


I read some things about PEP 612 here:

These seem related, but the examples given there still always copy-paste at least some part of the types.

How can this be done in pyright/mypy/general?

1

There are 1 answers

5
dROOOze On BEST ANSWER

You could simply pass [the function you want to copy the signature of] into [a decorator factory] which produces a no-op decorator that affects the typing API of the decorated function.

The following example can be checked on pyright-play.net (requires Python >= 3.12, as it uses syntax from PEP 695).

from __future__ import annotations
import typing_extensions as t

if t.TYPE_CHECKING:
    import collections.abc as cx

def withParameterAndReturnTypesOf[F: cx.Callable[..., t.Any]](f: F, /) -> cx.Callable[[F], F]:

    """
    Capture the exact type of `f`, then pretend that the decorated function is of this
    exact type.
    """

    return lambda _: _

def withParameterTypesOf[**P, R](f: cx.Callable[P, t.Any], /) -> cx.Callable[[cx.Callable[P, R]], cx.Callable[P, R]]:

    """
    Capture the parameters type of `f`, then pretend that the decorated function's
    parameters are of this type.

    `f`'s return type is ignored, and the decorated function's return type is preserved.
    """

    return lambda _: _
import os
from contextlib import contextmanager

@withParameterAndReturnTypesOf(os.open)
def my_open(*args: t.Any, **kwargs: t.Any) -> t.Any:
    return os.open(*args, **kwargs)

@contextmanager
@withParameterTypesOf(os.open)
def scoped_open(*args: t.Any, **kwargs: t.Any) -> cx.Generator[int, None, None]:
    fd = os.open(*args, **kwargs)
    try:
        yield fd
    finally:
        os.close(fd)
Expression of type "int" cannot be assigned to declared type "str"
  "int" is incompatible with "str"  (reportGeneralTypeIssues)
    vvvvvv   vvvvvvv
>>> a: str = my_open(1, 2)
                     ^
Argument of type "Literal[1]" cannot be assigned to parameter "path" of type "StrOrBytesPath" in function "open"
  Type "Literal[1]" cannot be assigned to type "StrOrBytesPath"
    "Literal[1]" is incompatible with "str"
    "Literal[1]" is incompatible with "bytes"
    "Literal[1]" is incompatible with protocol "PathLike[str]"
      "__fspath__" is not present
    "Literal[1]" is incompatible with protocol "PathLike[bytes]"
      "__fspath__" is not present  (reportGeneralTypeIssues)
                     v
>>> with scoped_open(1, 2) as a: pass
         ^^^^^^^^^^^       ^^^^
Expression of type "int" cannot be assigned to declared type "str"
  "int" is incompatible with "str"  (reportGeneralTypeIssues)