ENH: Add support for dtype string aliases to Series#astype

[skatsuta]

• Closes #xxxx (Replace xxxx with the Github issue number)
• Tests added: Please use assert_type() to assert the type of any return value

Problem

Series#astype has been improved in #519 to allow specific types to be inferred based on a value of the dtype argument.
I really appreciate the hard work that has been put into it.

However, that PR lacked support for aliases that specify dtype as a string.
Therefore, starting with version v1.5.3.230227, type checkers now report the following errors for code like below:

Sample code

s = pd.Series([1, 2, 3])
s_nullable_int = s.astype("Int64")
s_nullable_float = s.astype("Float64")

Errors reported by Pyright

error: Argument of type "Literal['Int64']" cannot be assigned to parameter "dtype" of type "ExtensionDtype" in function "astype"
    "Literal['Int64']" is incompatible with "ExtensionDtype" (reportGeneralTypeIssues)
error: Argument of type "Literal['Float64']" cannot be assigned to parameter "dtype" of type "ExtensionDtype" in function "astype"
    "Literal['Float64']" is incompatible with "ExtensionDtype" (reportGeneralTypeIssues)

Specifying dtype as a string is perfectly fine in pandas, so we need to fix the stub so that the errors are not reported for code like the above.

Solution

Add support for various dtype string aliases to (Boolean|Int|Str|Float|Complex)DtypeArg types, and add tests for passing each of those aliases as arguments to Series#astype.

Along with this, I have also refactored the stubs and tests to group code for the same dtype together so that it is easier to see what kind of dtype arguments are supported.

[Dr-Irv]

Thanks @skatsuta . I had a feeling that someone would pick up many of the possibilities that were missed when we put this together.

[skatsuta]

@Dr-Irv Thanks for merging quickly! Your hunch was right on the money ;)

Will this fix be released in a few days?
Since this issue is a sort of regression, and likely to have some impact on existing projects, including the one I'm working on, it might be good to release this patch relatively sooner than usual.

[Dr-Irv]

I should be able to do a release by Friday 3/3, if not earlier.

[zmoon]

Looks like 'category' is still missing?

[Dr-Irv]

Looks like 'category' is still missing?

No, it's there.

pandas-stubs/pandas-stubs/_typing.pyi

Line 151 in d2faa7f

CategoryDtypeArg: TypeAlias = Literal["category"]

[zmoon]

Weird, using v1.5.3.230227, which has that, I got this message

error: Argument 1 to "astype" of "DataFrame" has incompatible type "Literal['category']"; expected "Union[Union[Union[Type[bool], Type[bool_], BooleanDtype, Literal['bool']], Union[Type[int], Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype, Type[signedinteger[_8Bit]], Type[signedinteger[_16Bit]], Type[signedinteger[_32Bit]], Type[signedinteger[_64Bit]], Type[unsignedinteger[_8Bit]], Type[unsignedinteger[_16Bit]], Type[unsignedinteger[_32Bit]], Type[unsignedinteger[_64Bit]], Type[signedinteger[Any]], Type[unsignedinteger[Any]], Type[signedinteger[Any]], Type[unsignedinteger[Any]], Literal['int', 'int32']], Union[Type[str], StringDtype, Literal['str']], Type[bytes], Union[Float32Dtype, Float64Dtype, Type[floating[_16Bit]], Type[floating[_32Bit]], Type[floating[_64Bit]], Type[float], Literal['float']], Union[Type[complexfloating[_32Bit, _32Bit]], Type[complexfloating[_64Bit, _64Bit]], Type[complex], Literal['complex']], CategoricalDtype, ExtensionDtype, Literal['timedelta64[ns]', 'datetime64[ns]']], Mapping[Any, Union[ExtensionDtype, Union[str, dtype[generic], Type[str], Type[complex], Type[bool], Type[object]]]], Series[Any]]" [arg-type]

I guess because CategoryDtypeArg is missing from AstypeArg.

pandas-stubs/pandas-stubs/_typing.pyi

Lines 153 to 164 in d2faa7f

	AstypeArg: TypeAlias = (
	BooleanDtypeArg
	| IntDtypeArg
	| StrDtypeArg
	| BytesDtypeArg
	| FloatDtypeArg
	| ComplexDtypeArg
	| TimedeltaDtypeArg
	| TimestampDtypeArg
	| CategoricalDtype
	| ExtensionDtype
	)

[Dr-Irv]

I guess because CategoryDtypeArg is missing from AstypeArg.

Yes, you are correct. Can you open up a new issue for DataFrame.astype("category") and maybe follow up with a PR?