pandas-stubs icon indicating copy to clipboard operation
pandas-stubs copied to clipboard

Subclassing pandas DataFrame results in mypy misinterpreting indexers return

Open Rhapso10 opened this issue 1 year ago • 1 comments

Describe the bug

When subclassing pandas DataFrame by subclassing _constructor, type hints indicate that any return from a loc, iloc, etc... indexing is a pandas DataFrame not the subclass. To Reproduce

  1. Provide a minimal runnable pandas example that is not properly checked by the stubs.

import pandas as pd 


class MyClass(pd.DataFrame):
    @property
    def _constructor(self) -> 'type[MyClass]':
        return MyClass

df = MyClass({'a': [1, 2, 3]})
df2 = df.iloc[1:2]
reveal_type(df2)

  1. Indicate which type checker you are using (mypy or pyright).

I used mypy.

  1. Show the error message received from that type checker while checking your example. image
  • OS: Linux
  • OS Version: Debian GNU/Linux 10
  • python version: 3.10.13
  • version of type checker : mypy==1.8.0
  • version of installed pandas-stubs: 2.1.4.231227

Rhapso10 avatar Jan 23 '24 15:01 Rhapso10

We should probably return Self more often: after double checking that pandas actually returns the sub-class, for example, _consructor returns DataFrame unless a sub-class overwrites it.

twoertwein avatar Jan 23 '24 15:01 twoertwein