WIP

Author: avinashpancham

pandas/_typing.py

@@ -31,7 +31,7 @@
 
     from pandas._libs import Period, Timedelta, Timestamp
 
-    from pandas.core.dtypes.dtypes import ExtensionDtype
+    from pandas.core.dtypes.dtypes import DatetimeTZDtype, ExtensionDtype, IntervalDtype
 
     from pandas import Interval
     from pandas.core.arrays.base import ExtensionArray  # noqa: F401
@@ -95,7 +95,8 @@
 
 # dtypes
 NpDtype = Union[str, np.dtype, Type[Union[str, float, int, complex, bool, object]]]
-Dtype = Union["ExtensionDtype", NpDtype]
+PandasDtype = Union[IntervalDtype, DatetimeTZDtype]
+Dtype = Union["ExtensionDtype", NpDtype, PandasDtype]
 # DtypeArg specifies all allowable dtypes in a functions its dtype argument
 DtypeArg = Union[Dtype, Dict[Label, Dtype]]
 DtypeObj = Union[np.dtype, "ExtensionDtype"]

pandas/core/arrays/base.py

@@ -189,7 +189,7 @@ class ExtensionArray:
     # ------------------------------------------------------------------------
 
     @classmethod
-    def _from_sequence(cls, scalars, *, dtype: Optional[Dtype] = None, copy=False):
+    def _from_sequence(cls, scalars, *, dtype=None, copy=False):
         """
         Construct a new ExtensionArray from a sequence of scalars.
 
@@ -211,9 +211,7 @@ def _from_sequence(cls, scalars, *, dtype: Optional[Dtype] = None, copy=False):
         raise AbstractMethodError(cls)
 
     @classmethod
-    def _from_sequence_of_strings(
-        cls, strings, *, dtype: Optional[Dtype] = None, copy=False
-    ):
+    def _from_sequence_of_strings(cls, strings, *, dtype=None, copy=False):
         """
         Construct a new ExtensionArray from a sequence of strings.
 

pandas/core/arrays/categorical.py

@@ -303,7 +303,7 @@ def __init__(
         values,
         categories=None,
         ordered=None,
-        dtype: Optional[Dtype] = None,
+        dtype: Optional["CategoricalDtype"] = None,
         fastpath=False,
         copy: bool = True,
     ):
@@ -405,7 +405,9 @@ def _constructor(self) -> Type["Categorical"]:
         return Categorical
 
     @classmethod
-    def _from_sequence(cls, scalars, *, dtype: Optional[Dtype] = None, copy=False):
+    def _from_sequence(
+        cls, scalars, *, dtype: Optional["CategoricalDtype"] = None, copy=False
+    ):
         return Categorical(scalars, dtype=dtype, copy=copy)
 
     def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
@@ -541,7 +543,11 @@ def _from_inferred_categories(
 
     @classmethod
     def from_codes(
-        cls, codes, categories=None, ordered=None, dtype: Optional[Dtype] = None
+        cls,
+        codes,
+        categories=None,
+        ordered=None,
+        dtype: Optional["CategoricalDtype"] = None,
     ):
         """
         Make a Categorical type from codes and categories or dtype.

pandas/core/arrays/datetimes.py

@@ -218,7 +218,11 @@ class DatetimeArray(dtl.TimelikeOps, dtl.DatelikeOps):
     _freq = None
 
     def __init__(
-        self, values, dtype: Optional[Dtype] = DT64NS_DTYPE, freq=None, copy=False
+        self,
+        values,
+        dtype: Optional[DatetimeTZDtype] = DT64NS_DTYPE,
+        freq=None,
+        copy=False,
     ):
         if isinstance(values, (ABCSeries, ABCIndex)):
             values = values._values

pandas/core/arrays/interval.py

@@ -170,7 +170,7 @@ def __new__(
         cls,
         data,
         closed=None,
-        dtype: Optional[Dtype] = None,
+        dtype: Optional[IntervalDtype] = None,
         copy: bool = False,
         verify_integrity: bool = True,
     ):
@@ -217,7 +217,7 @@ def _simple_new(
         right,
         closed=None,
         copy=False,
-        dtype: Optional[Dtype] = None,
+        dtype: Optional[IntervalDtype] = None,
         verify_integrity=True,
     ):
         result = IntervalMixin.__new__(cls)

pandas/core/arrays/period.py

@@ -223,7 +223,7 @@ def _from_sequence(
 
     @classmethod
     def _from_sequence_of_strings(
-        cls, strings, *, dtype: Optional[Dtype] = None, copy=False
+        cls, strings, *, dtype: Optional[PeriodDtype] = None, copy=False
     ) -> "PeriodArray":
         return cls._from_sequence(strings, dtype=dtype, copy=copy)
 

pandas/core/indexes/category.py

@@ -7,14 +7,15 @@
 
 from pandas._libs import index as libindex
 from pandas._libs.lib import no_default
-from pandas._typing import ArrayLike, Dtype, Label
+from pandas._typing import ArrayLike, Label
 from pandas.util._decorators import Appender, doc
 
 from pandas.core.dtypes.common import (
     ensure_platform_int,
     is_categorical_dtype,
     is_scalar,
 )
+from pandas.core.dtypes.dtypes import CategoricalDtype
 from pandas.core.dtypes.missing import is_valid_nat_for_dtype, isna, notna
 
 from pandas.core import accessor
@@ -184,7 +185,7 @@ def __new__(
         data=None,
         categories=None,
         ordered=None,
-        dtype: Optional[Dtype] = None,
+        dtype: Optional["CategoricalDtype"] = None,
         copy=False,
         name=None,
     ):

pandas/core/internals/concat.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from pandas._libs import NaT, internals as libinternals
-from pandas._typing import DtypeObj, Shape
+from pandas._typing import Dtype, DtypeObj, Shape
 from pandas.util._decorators import cache_readonly
 
 from pandas.core.dtypes.cast import maybe_promote
@@ -235,7 +235,7 @@ def is_na(self) -> bool:
 
         return isna_all(values_flat)
 
-    def get_reindexed_values(self, empty_dtype: DtypeObj, upcasted_na):
+    def get_reindexed_values(self, empty_dtype: Dtype, upcasted_na):
         if upcasted_na is None:
             # No upcasting is necessary
             fill_value = self.block.fill_value

pandas/core/reshape/tile.py

@@ -7,7 +7,6 @@
 
 from pandas._libs import Timedelta, Timestamp
 from pandas._libs.lib import infer_dtype
-from pandas._typing import Dtype
 
 from pandas.core.dtypes.common import (
     DT64NS_DTYPE,
@@ -27,7 +26,14 @@
 from pandas.core.dtypes.generic import ABCSeries
 from pandas.core.dtypes.missing import isna
 
-from pandas import Categorical, Index, IntervalIndex, to_datetime, to_timedelta
+from pandas import (
+    Categorical,
+    DatetimeTZDtype,
+    Index,
+    IntervalIndex,
+    to_datetime,
+    to_timedelta,
+)
 import pandas.core.algorithms as algos
 import pandas.core.nanops as nanops
 
@@ -378,7 +384,7 @@ def _bins_to_cuts(
     labels=None,
     precision: int = 3,
     include_lowest: bool = False,
-    dtype: Optional[Dtype] = None,
+    dtype: Optional[DatetimeTZDtype] = None,
     duplicates: str = "raise",
     ordered: bool = True,
 ):
@@ -543,7 +549,7 @@ def _format_labels(
     precision: int,
     right: bool = True,
     include_lowest: bool = False,
-    dtype: Optional[Dtype] = None,
+    dtype: Optional[DatetimeTZDtype] = None,
 ):
     """ based on the dtype, return our labels """
     closed = "right" if right else "left"

pandas/core/series.py

@@ -1309,7 +1309,7 @@ def to_string(
         header=True,
         index=True,
         length=False,
-        dtype: bool = False,
+        dtype: Dtype = False,
         name=False,
         max_rows=None,
         min_rows=None,

pandas/io/formats/format.py

@@ -42,6 +42,7 @@
     ColspaceArgType,
     ColspaceType,
     CompressionOptions,
+    Dtype,
     FilePathOrBuffer,
     FloatFormatType,
     FormattersType,
@@ -244,7 +245,7 @@ def __init__(
         na_rep: str = "NaN",
         name: bool = False,
         float_format: Optional[str] = None,
-        dtype: bool = True,
+        dtype: Dtype = True,
         max_rows: Optional[int] = None,
         min_rows: Optional[int] = None,
     ):

pandas/io/json/_json.py

@@ -12,6 +12,7 @@
 from pandas._typing import (
     CompressionOptions,
     DtypeArg,
+    FrameOrSeriesUnion,
     IndexLabel,
     JSONSerializable,
     StorageOptions,
@@ -810,7 +811,7 @@ def __init__(
         self.convert_dates = convert_dates
         self.date_unit = date_unit
         self.keep_default_dates = keep_default_dates
-        self.obj = None
+        self.obj: Optional[FrameOrSeriesUnion] = None
 
     def check_keys_split(self, decoded):
         """

pandas/io/sql.py

@@ -1110,7 +1110,7 @@ def _harmonize_columns(self, parse_dates=None):
 
     def _sqlalchemy_type(self, col):
 
-        dtype = self.dtype or {}
+        dtype: Union[DtypeArg, Dict] = self.dtype or {}
         if col.name in dtype:
             return self.dtype[col.name]
 
@@ -1736,7 +1736,7 @@ def _create_table_setup(self):
         return create_stmts
 
     def _sql_type_name(self, col):
-        dtype = self.dtype or {}
+        dtype: Union[DtypeArg, Dict] = self.dtype or {}
         if col.name in dtype:
             return dtype[col.name]