Skip to content

Boolean indexing doesn't work with subclass and TypeVar #1069

Closed
@davetapley

Description

@davetapley

Describe the bug
A suggested on #908 (comment) I'm trying to use DataFrameT = TypeVar("DataFrameT", bound=DataFrame),
but with boolean indexing instead of a pipe.

To Reproduce

from typing import TypeVar, reveal_type

from pandas import DataFrame, Series


class SubDF(DataFrame):
    # https://pandas.pydata.org/pandas-docs/stable/development/extending.html#override-constructor-properties
    @property
    def _constructor(self):
        return SubDF

    @property
    def _constructor_sliced(self):
        return Series


sub = SubDF({'a': [1, 2, 3]})

DataFrameT = TypeVar("DataFrameT", bound=DataFrame)


def func(df: DataFrameT) -> DataFrameT:
    index = Series([True, False, True])

    df_ = df.loc[index]
    reveal_type(df_)

    return df_  # Type "DataFrame" is not assignable to return type "DataFrameT@func"


reveal_type(func(sub))

pyright:

  /workspaces/ng/repro.py:27:17 - information: Type of "df_" is "DataFrame"
  /workspaces/ng/repro.py:29:12 - error: Type "DataFrame" is not assignable to return type "DataFrameT@func"
    Type "DataFrame" is not assignable to type "DataFrameT@func" (reportReturnType)
  /workspaces/ng/repro.py:32:13 - information: Type of "func(sub)" is "SubDF"
1 error, 0 warnings, 2 informations 

Please complete the following information:

  • OS: Linux
  • OS Version 20.04.6
  • python version 3.12.2
  • version of type checker 1.1.390
  • version of installed pandas-stubs 2.2.3.241126

Additional context

Repro inspired by:

Might be same root cause as:

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions