fix(typing): #983 return type of StringMethods.match

[cmp0xff]

• Closes df.columns.str.match actually gives npt.NDArray[np.bool_], but mypy thinks it is pd.Index[str] #983
• Tests added: Please use assert_type() to assert the type of any return value

[cmp0xff]

Hi @Dr-Irv thank you for the comment. In #983 (comment) you mentioned:

The declaration of match() in core/strings.pyi is incorrect.

I partially fixed it.

• Now the typing is correct and more accurate for Series.str.match.
• However, the return type actually depends on where str.match is called.
  • If it is from a Series, the return type is Series[bool], as is fixed
  • If it is from an Index, the return type is npt.NDArray[np.bool_], as in df.columns.str.match actually gives npt.NDArray[np.bool_], but mypy thinks it is pd.Index[str] #983

How can I capture this expectation?

But to fix it, the StringMethods class will need an additional argument to pass in the expected result of match(), similar to what is done with str.split().

I do not get this suggestion. Could you elaborate?

PR with tests welcome

#990

[Dr-Irv]

How can I capture this expectation?

But to fix it, the StringMethods class will need an additional argument to pass in the expected result of match(), similar to what is done with str.split().

I do not get this suggestion. Could you elaborate?

In core/strings.pyi, we have this:

# The _TS type is what is used for the result of str.split with expand=True
_TS = TypeVar("_TS", DataFrame, MultiIndex)

class StringMethods(NoNewAttributesMixin, Generic[T, _TS]):

You need to create an additional TypeVar, let's call it _MR (for match result). Then the Generic part of StringMethods adds _MR to the parameters of Generic.

Then the match() method would return _MR.

Then you change the references to StringMethods in core/series.pyi and core/indexes/base.pyi to pass in the correct type of result you want match to return as the 3rd argument.

See the pattern used with split() - that's how we control the return type of that particular method.

I'll let you address that change before firing off the CI.

[cmp0xff]

Hi @Dr-Irv , thanks for the explanations. I have made the changes.

[Dr-Irv]

In the stubs, instead of using npt.NDArray[np.bool_], can you use np_ndarray_bool which you can import from pandas._typing .

[Dr-Irv]

Thanks @cmp0xff