Add support for DataFrame#loc[..., NDArray[bool]]

[skatsuta]

• Tests added: Please use assert_type() to assert the type of any return value

Problem

Currently, only list[bool] or Series[bool] can be used as a boolean vector for column masking in DataFrame#loc.
For example, Pyright reports the following error when using NDArray[bool] for column masking:

# test.py
import pandas as pd

df = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
df.loc[:, df.columns.isin(["col1"])]

$ pyright test.py
/Users/skatsuta/src/github.com/SmartDriveInc/mobility-risk-score/ml/test.py
  /Users/skatsuta/src/github.com/SmartDriveInc/mobility-risk-score/ml/test.py:4:1 - error: No overloads for "__getitem__" match the provided arguments (reportCallIssue)
  /Users/skatsuta/src/github.com/SmartDriveInc/mobility-risk-score/ml/test.py:4:1 - error: Argument of type "tuple[slice, np_ndarray_bool]" cannot be assigned to parameter "idx" of type "tuple[Scalar, slice]" in function "__getitem__"
    "tuple[slice, np_ndarray_bool]" is incompatible with "tuple[Scalar, slice]"
      Tuple entry 1 is incorrect type
        Type "slice" cannot be assigned to type "Scalar"
          "slice" is incompatible with "str"
          "slice" is incompatible with "bytes"
          "slice" is incompatible with "date"
          "slice" is incompatible with "datetime"
          "slice" is incompatible with "timedelta" (reportArgumentType)
2 errors, 0 warnings, 0 informations

However, such usage is perfectly valid from a pandas perspective.

Solution

Allow MaskType to be used as a boolean vector for column masking in DataFrame#loc.
MaskType contains np_ndarray_bool in addition to Series[bool], so the above code will no longer generate a type error.

Instead, Series[bool] is removed from the type annotation because it is included in MaskType.

Also add test cases for it.

[Dr-Irv]

Thanks @skatsuta

I created an issue #868 to see if we can do a similar change elsewhere.