diff --git a/pandas/core/arrays/integer.py b/pandas/core/arrays/integer.py index 022e6a7322872..2c60cc4387158 100644 --- a/pandas/core/arrays/integer.py +++ b/pandas/core/arrays/integer.py @@ -1,5 +1,5 @@ import numbers -from typing import Any, Tuple, Type +from typing import TYPE_CHECKING, Any, Dict, Tuple, Type, Union import warnings import numpy as np @@ -31,6 +31,9 @@ from .masked import BaseMaskedArray +if TYPE_CHECKING: + import pyarrow # noqa: F401 + class _IntegerDtype(ExtensionDtype): """ @@ -52,33 +55,33 @@ def __repr__(self) -> str: return f"{sign}Int{8 * self.itemsize}Dtype()" @cache_readonly - def is_signed_integer(self): + def is_signed_integer(self) -> bool: return self.kind == "i" @cache_readonly - def is_unsigned_integer(self): + def is_unsigned_integer(self) -> bool: return self.kind == "u" @property - def _is_numeric(self): + def _is_numeric(self) -> bool: return True @cache_readonly - def numpy_dtype(self): + def numpy_dtype(self) -> np.dtype: """ Return an instance of our numpy dtype """ return np.dtype(self.type) @cache_readonly - def kind(self): + def kind(self) -> str: return self.numpy_dtype.kind @cache_readonly - def itemsize(self): + def itemsize(self) -> int: """ Return the number of bytes in this dtype """ return self.numpy_dtype.itemsize @classmethod - def construct_array_type(cls): + def construct_array_type(cls) -> Type["IntegerArray"]: """ Return the array type associated with this dtype. @@ -88,9 +91,13 @@ def construct_array_type(cls): """ return IntegerArray - def __from_arrow__(self, array): - """Construct IntegerArray from passed pyarrow Array/ChunkedArray""" - import pyarrow + def __from_arrow__( + self, array: Union["pyarrow.Array", "pyarrow.ChunkedArray"] + ) -> "IntegerArray": + """ + Construct IntegerArray from pyarrow Array/ChunkedArray. + """ + import pyarrow # noqa: F811 from pandas.core.arrays._arrow_utils import pyarrow_array_to_numpy_and_mask if isinstance(array, pyarrow.Array): @@ -108,7 +115,7 @@ def __from_arrow__(self, array): return IntegerArray._concat_same_type(results) -def integer_array(values, dtype=None, copy=False): +def integer_array(values, dtype=None, copy: bool = False,) -> "IntegerArray": """ Infer and return an integer array of the values. @@ -131,7 +138,7 @@ def integer_array(values, dtype=None, copy=False): return IntegerArray(values, mask) -def safe_cast(values, dtype, copy): +def safe_cast(values, dtype, copy: bool): """ Safely cast the values to the dtype if they are equivalent, meaning floats must be equivalent to the @@ -152,7 +159,9 @@ def safe_cast(values, dtype, copy): ) -def coerce_to_array(values, dtype, mask=None, copy=False): +def coerce_to_array( + values, dtype, mask=None, copy: bool = False, +) -> Tuple[np.ndarray, np.ndarray]: """ Coerce the input values array to numpy arrays with a mask @@ -322,10 +331,10 @@ class IntegerArray(BaseMaskedArray): _internal_fill_value = 1 @cache_readonly - def dtype(self): + def dtype(self) -> _IntegerDtype: return _dtypes[str(self._data.dtype)] - def __init__(self, values, mask, copy=False): + def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False): if not (isinstance(values, np.ndarray) and is_integer_dtype(values.dtype)): raise TypeError( "values should be integer numpy array. Use " @@ -345,21 +354,23 @@ def __init__(self, values, mask, copy=False): self._mask = mask @classmethod - def _from_sequence(cls, scalars, dtype=None, copy=False): + def _from_sequence(cls, scalars, dtype=None, copy: bool = False) -> "IntegerArray": return integer_array(scalars, dtype=dtype, copy=copy) @classmethod - def _from_sequence_of_strings(cls, strings, dtype=None, copy=False): + def _from_sequence_of_strings( + cls, strings, dtype=None, copy: bool = False + ) -> "IntegerArray": scalars = to_numeric(strings, errors="raise") return cls._from_sequence(scalars, dtype, copy) @classmethod - def _from_factorized(cls, values, original): + def _from_factorized(cls, values, original) -> "IntegerArray": return integer_array(values, dtype=original.dtype) _HANDLED_TYPES = (np.ndarray, numbers.Number) - def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + def __array_ufunc__(self, ufunc, method: str, *inputs, **kwargs): # For IntegerArray inputs, we apply the ufunc to ._data # and mask the result. if method == "reduce": @@ -697,103 +708,65 @@ def integer_arithmetic_method(self, other): """ # create the Dtype -Int8Dtype = register_extension_dtype( - type( - "Int8Dtype", - (_IntegerDtype,), - { - "type": np.int8, - "name": "Int8", - "__doc__": _dtype_docstring.format(dtype="int8"), - }, - ) -) -Int16Dtype = register_extension_dtype( - type( - "Int16Dtype", - (_IntegerDtype,), - { - "type": np.int16, - "name": "Int16", - "__doc__": _dtype_docstring.format(dtype="int16"), - }, - ) -) -Int32Dtype = register_extension_dtype( - type( - "Int32Dtype", - (_IntegerDtype,), - { - "type": np.int32, - "name": "Int32", - "__doc__": _dtype_docstring.format(dtype="int32"), - }, - ) -) +@register_extension_dtype +class Int8Dtype(_IntegerDtype): + type = np.int8 + name = "Int8" + __doc__ = _dtype_docstring.format(dtype="int8") -Int64Dtype = register_extension_dtype( - type( - "Int64Dtype", - (_IntegerDtype,), - { - "type": np.int64, - "name": "Int64", - "__doc__": _dtype_docstring.format(dtype="int64"), - }, - ) -) -UInt8Dtype = register_extension_dtype( - type( - "UInt8Dtype", - (_IntegerDtype,), - { - "type": np.uint8, - "name": "UInt8", - "__doc__": _dtype_docstring.format(dtype="uint8"), - }, - ) -) +@register_extension_dtype +class Int16Dtype(_IntegerDtype): + type = np.int16 + name = "Int16" + __doc__ = _dtype_docstring.format(dtype="int16") -UInt16Dtype = register_extension_dtype( - type( - "UInt16Dtype", - (_IntegerDtype,), - { - "type": np.uint16, - "name": "UInt16", - "__doc__": _dtype_docstring.format(dtype="uint16"), - }, - ) -) -UInt32Dtype = register_extension_dtype( - type( - "UInt32Dtype", - (_IntegerDtype,), - { - "type": np.uint32, - "name": "UInt32", - "__doc__": _dtype_docstring.format(dtype="uint32"), - }, - ) -) +@register_extension_dtype +class Int32Dtype(_IntegerDtype): + type = np.int32 + name = "Int32" + __doc__ = _dtype_docstring.format(dtype="int32") + + +@register_extension_dtype +class Int64Dtype(_IntegerDtype): + type = np.int64 + name = "Int64" + __doc__ = _dtype_docstring.format(dtype="int64") + + +@register_extension_dtype +class UInt8Dtype(_IntegerDtype): + type = np.uint8 + name = "UInt8" + __doc__ = _dtype_docstring.format(dtype="uint8") + + +@register_extension_dtype +class UInt16Dtype(_IntegerDtype): + type = np.uint16 + name = "UInt16" + __doc__ = _dtype_docstring.format(dtype="uint16") + + +@register_extension_dtype +class UInt32Dtype(_IntegerDtype): + type = np.uint32 + name = "UInt32" + __doc__ = _dtype_docstring.format(dtype="uint32") + + +@register_extension_dtype +class UInt64Dtype(_IntegerDtype): + type = np.uint64 + name = "UInt64" + __doc__ = _dtype_docstring.format(dtype="uint64") -UInt64Dtype = register_extension_dtype( - type( - "UInt64Dtype", - (_IntegerDtype,), - { - "type": np.uint64, - "name": "UInt64", - "__doc__": _dtype_docstring.format(dtype="uint64"), - }, - ) -) -_dtypes = { +_dtypes: Dict[str, _IntegerDtype] = { "int8": Int8Dtype(), "int16": Int16Dtype(), "int32": Int32Dtype(),