Fix types of declared ExtensionDtypes (#329)

Author: Dr-Irv

pandas-stubs/__init__.pyi

@@ -25,6 +25,8 @@ from .core.api import (
     DateOffset as DateOffset,
     DatetimeIndex as DatetimeIndex,
     DatetimeTZDtype as DatetimeTZDtype,
+    Float32Dtype as Float32Dtype,
+    Float64Dtype as Float64Dtype,
     Float64Index as Float64Index,
     Grouper as Grouper,
     Index as Index,

pandas-stubs/core/api.pyi

@@ -5,6 +5,10 @@ from pandas.core.algorithms import (
 )
 from pandas.core.arrays import Categorical as Categorical
 from pandas.core.arrays.boolean import BooleanDtype as BooleanDtype
+from pandas.core.arrays.floating import (
+    Float32Dtype as Float32Dtype,
+    Float64Dtype as Float64Dtype,
+)
 from pandas.core.arrays.integer import (
     Int8Dtype as Int8Dtype,
     Int16Dtype as Int16Dtype,

pandas-stubs/core/arrays/floating.pyi

@@ -0,0 +1,4 @@
+from pandas.core.arrays.numeric import NumericDtype
+
+class Float32Dtype(NumericDtype): ...
+class Float64Dtype(NumericDtype): ...

pandas-stubs/core/arrays/integer.pyi

@@ -28,11 +28,11 @@ class IntegerArray(BaseMaskedArray):
     def __setitem__(self, key, value) -> None: ...
     def astype(self, dtype, copy: bool = ...): ...
 
-class Int8Dtype: ...
-class Int16Dtype: ...
-class Int32Dtype: ...
-class Int64Dtype: ...
-class UInt8Dtype: ...
-class UInt16Dtype: ...
-class UInt32Dtype: ...
-class UInt64Dtype: ...
+class Int8Dtype(_IntegerDtype): ...
+class Int16Dtype(_IntegerDtype): ...
+class Int32Dtype(_IntegerDtype): ...
+class Int64Dtype(_IntegerDtype): ...
+class UInt8Dtype(_IntegerDtype): ...
+class UInt16Dtype(_IntegerDtype): ...
+class UInt32Dtype(_IntegerDtype): ...
+class UInt64Dtype(_IntegerDtype): ...

pandas-stubs/core/arrays/numeric.pyi

@@ -0,0 +1,3 @@
+from pandas.core.dtypes.dtypes import BaseMaskedDtype
+
+class NumericDtype(BaseMaskedDtype): ...

pandas-stubs/core/dtypes/dtypes.pyi

@@ -17,6 +17,8 @@ _str = str
 
 def register_extension_dtype(cls: type[ExtensionDtype]) -> type[ExtensionDtype]: ...
 
+class BaseMaskedDtype(ExtensionDtype): ...
+
 class PandasExtensionDtype(ExtensionDtype):
     subdtype = ...
     str: _str | None = ...

tests/test_api_types.py

@@ -1,5 +1,8 @@
+from typing import Type
+
 import numpy as np
 import pandas as pd
+from pandas.api.extensions import ExtensionDtype
 import pandas.api.types as api
 from typing_extensions import assert_type
 
@@ -176,3 +179,28 @@ def test_infer_dtype() -> None:
 def test_union_categoricals() -> None:
     to_union = [pd.Categorical([1, 2, 3]), pd.Categorical([3, 4, 5])]
     check(assert_type(api.union_categoricals(to_union), pd.Categorical), pd.Categorical)
+
+
+def check_extension_dtypes() -> None:
+    # GH 315
+    def check_ext_dtype(etype: Type[ExtensionDtype]):
+        assert issubclass(etype, ExtensionDtype)
+
+    check_ext_dtype(pd.Int64Dtype)
+    check_ext_dtype(pd.Int8Dtype)
+    check_ext_dtype(pd.Int16Dtype)
+    check_ext_dtype(pd.Int32Dtype)
+    check_ext_dtype(pd.Int64Dtype)
+    check_ext_dtype(pd.UInt8Dtype)
+    check_ext_dtype(pd.UInt16Dtype)
+    check_ext_dtype(pd.UInt32Dtype)
+    check_ext_dtype(pd.UInt64Dtype)
+    check_ext_dtype(pd.BooleanDtype)
+    check_ext_dtype(pd.StringDtype)
+    check_ext_dtype(pd.CategoricalDtype)
+    check_ext_dtype(pd.DatetimeTZDtype)
+    check_ext_dtype(pd.IntervalDtype)
+    check_ext_dtype(pd.PeriodDtype)
+    check_ext_dtype(pd.SparseDtype)
+    check_ext_dtype(pd.Float32Dtype)
+    check_ext_dtype(pd.Float64Dtype)