How can I distinguish between two dynamic types?

Viewed 251

I am type checking a project that makes use of pandas. Pandas doesn't contain type annotations and there isn't a stub file in the typeshed.

As I would expect, mypy raises an error for this simple example:

class A:
    def method(self) -> int:
        return 1


class B(A):
    def method(self) -> float:
        return 1.1
$ mypy mypy_example.py 
mypy_example.py:11: error: Return type "float" of "method" incompatible with return type "int" in supertype "A"

Consider the following example:

class C:
    def method(self) -> pd.Series:
        return pd.Series([1, 2, 3])


class D(C):
    def method(self) -> pd.DataFrame:
        return pd.DataFrame({"a": [1, 2, 3]})

As expected, mypy says no stub file can be found for pandas, so it doesn't find the error.

$ mypy mypy_example.py 
mypy_example.py:1: error: No library stub file for module 'pandas'
mypy_example.py:1: note: (Stub files are from https://github.com/python/typeshed)
mypy_example.py:11: error: Return type "float" of "method" incompatible with return type "int" in supertype "A"

I could set ignore_missing_imports, but that means I miss the error I want to catch.

I've tried a few things in stub files without success:

from typing import Any, NewType

# dynamic typing, but doesn't discriminate between Series and DataFrame
Series = Any
DataFrame = Any

# discriminates but doesn't dynamically type
Series = NewType('Series', object)
DataFrame = NewType('DataFrame', object)

Is it possible to write a short stub file or type annotation that will allow me to take advantage of dynamic-typing but recognise that pd.Series and pd.DataFrame are different types?

1 Answers

Rather then trying to get mypy to distinguish between two dynamic classes, I would instead actually take the route of making them not dynamic (or rather, only partially dynamic) by defining them as full-fledged classes inside of your stubs.

You can get started with a very preliminary and minimal set of stubs by defining your two classes to look like this:

from typing import Any

class Series:
    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
    def __getattr__(self, name: str) -> Any: ...

class DataFrame(Series):
    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
    def __getattr__(self, name: str) -> Any: ...

The __getattr__ function lets mypy understand that your class is incomplete and not fully annotated. This means doing things like DataFrame().query(...) will continue to type-check, even though that function was never explicitly added to your class.

Of course, if you do decide to add in a few method signatures, mypy will begin typechecking those calls instead of leaving them dynamically typed. This means you can also progressively add in more precise method signatures if you want, and maybe eventually get rid of the __getattr__ entirely.

If you decide to take this route, you might find the existing numpy stubs to be a good source of inspiration. And if you want really precise types, the discussion here is probably relevant.

The typeshed guidelines on writing incomplete stubs have more information on writing partial stubs, if you're curious.

Related