I've written this code:

from typing import overload, TYPE_CHECKING, Protocol, Any

import pyarrow as pa  # type: ignore[import-not-found]

class PyArrowArray(Protocol):
    @property
    def buffers(self) -> Any: ...

@overload
def func(a: PyArrowArray) -> int: ...
@overload
def func(a: str) -> str: ...
@overload
def func(a: Any) -> str | int: ...


def func(a) -> str | int:
    if isinstance(a, pa.Array):
        return 0
    return '0'
    
reveal_type(func(pa.array([1,2,3])))

PyArrow is a Python library which does not have type hints. However, there is a package pyarrow-stubs which provides types for it.

I have a function can accept either a pyarrow.Array or a str:

• if it receives a pyarrow.Array, it returns an int
• if it receives a str, it returns a str

I would like to annotate it such that:

• if a user has pyarrow-stubs installed, then func(pa.array([1,2,3])) is revealed to be int
• if a user doesn't have pyarrow-stubs installed, then func(pa.array([1,2,3])) should be revealed to be int | str, because pa is not known statically

I was hoping that the code above would accomplish that, but it doesn't. If pyarrow-stubs is not installed, I get

Revealed type is "Any"

I was expecting that the

@overload
def func(a: Any) -> str | int: ...

overload would be matched and that I'd get Revealed type is int | str

I've written this code:

from typing import overload, TYPE_CHECKING, Protocol, Any

import pyarrow as pa  # type: ignore[import-not-found]

class PyArrowArray(Protocol):
    @property
    def buffers(self) -> Any: ...

@overload
def func(a: PyArrowArray) -> int: ...
@overload
def func(a: str) -> str: ...
@overload
def func(a: Any) -> str | int: ...


def func(a) -> str | int:
    if isinstance(a, pa.Array):
        return 0
    return '0'
    
reveal_type(func(pa.array([1,2,3])))

PyArrow is a Python library which does not have type hints. However, there is a package pyarrow-stubs which provides types for it.

I have a function can accept either a pyarrow.Array or a str:

• if it receives a pyarrow.Array, it returns an int
• if it receives a str, it returns a str

I would like to annotate it such that:

• if a user has pyarrow-stubs installed, then func(pa.array([1,2,3])) is revealed to be int
• if a user doesn't have pyarrow-stubs installed, then func(pa.array([1,2,3])) should be revealed to be int | str, because pa is not known statically

I was hoping that the code above would accomplish that, but it doesn't. If pyarrow-stubs is not installed, I get

Revealed type is "Any"

I was expecting that the

@overload
def func(a: Any) -> str | int: ...

overload would be matched and that I'd get Revealed type is int | str

• python
• python-typing
• mypy

• A more minimal case could be written with: bla = cast(Any, None); reveal_type(func(bla)) # Unknown/Any, independent of pyarrow, neither pyright nor mypy likes that. – Daraan Commented Mar 14 at 17:34
• If stubs are missing, the first overload is also Any and thus matches everything. Imports from untyped packages are just Any. Just make those stubs a dependency of your package (maybe in a [types]optional group that should only be installed in development) and don't try to make fallbacks if they are missing - python typing is rather lax even with all strictness flags enabled, there's really no point in supporting something in absence of typing info that can be trivially acquired. – STerliakov Commented Mar 14 at 18:26
• 1 @STerliakov if the first overload matches everything, shouldn't the revealed type be int? – dROOOze Commented Mar 14 at 21:12
• 1 @dROOOze good catch! My bad, there's a special case for Any in overloads spec breaking even that inference. If two or more overloads match due to Any presence and have incompatible return types, Any is inferred. – STerliakov Commented Mar 14 at 22:55
• Oh, okay - well anyway, it's safe to say that Python type-checkers aren't designed to handle the case in the question. mypy can be configured to behave like pyright using --follow-untyped-imports (or its equivalent configuration file option), but that doesn't help here because pyarrow's definition files are written in Cython. – dROOOze Commented Mar 14 at 23:00

1 Answer 1

I am confident to say that this is very likely not possible. The Any Type per PEP 484 matches everything, hence an Any input will match all overloads. What happens in such cases is defined here in the Mypy docs:

[...] if multiple variants match due to an argument being of type Any, mypy will make the inferred type also be Any:

pyright docs also describes it similarly, interestingly with pyright you can choose one, but not both overloads and avoid Unknown, this is explained here.

# NOTE: pyright only

@overload
def func(a: PyArrowArray) -> int: ...

@overload
def func(a: Any) -> str | int: ...

def func(a):
    if isinstance(a, pa.Array):
        return 0
    return "0"


bla = cast(PyArrowArray, ...)

reveal_type(func(bla))     # int
reveal_type(func("fooo"))  # str | int  :(
reveal_type(func(pa.array([1, 2, 3])))  # str | int  # Not Unknown

The only solution I somewhat see is that in turn pa.array is not allowed to be Any. Currently I see no good way to satisfy that without destroying compatibility when stubs are present. Best make it a requirement. Or somehow make everything a no-op class that is secondary to the stubs.

I assume you are looking for a mypy solution, with pyright you can solve it like this:

from typing import overload, TYPE_CHECKING, Callable, Protocol, Any, reveal_type, T
import pyarrow as pa

if TYPE_CHECKING:
    if pa.array is None:  # path only taken when stubs not present
        class unknown:

            def __call__(self, *args, **kwargs) -> unknown:
                ...

        pa.array = unknown()

# ... your original code
# type of pa.array without stubs will be: array: unknown | Unknown

reveal_type(func(bla)) # int
reveal_type(func("fooo")) # str
reveal_type(func(pa.array([1, 2, 3])))  # int | str