TYP: nattype.pyi (#40503)

jbrockmendel · web-flow · commit f841eda7fcdd · 2021-03-22T08:43:14.000-04:00
diff --git a/pandas/_libs/tslibs/nattype.pyi b/pandas/_libs/tslibs/nattype.pyi
@@ -0,0 +1,135 @@
+
+from datetime import datetime
+
+import numpy as np
+
+NaT: NaTType
+iNaT: int
+nat_strings: set[str]
+
+def is_null_datetimelike(val: object, inat_is_null: bool = ...) -> bool: ...
+
+class NaTType(datetime):
+    value: np.int64
+
+    def asm8(self) -> np.datetime64: ...
+    def to_datetime64(self) -> np.datetime64: ...
+    def to_numpy(self, dtype=..., copy: bool = ...) -> np.datetime64: ...
+
+    @property
+    def is_leap_year(self) -> bool: ...
+    @property
+    def is_month_start(self) -> bool: ...
+    @property
+    def is_quarter_start(self) -> bool: ...
+    @property
+    def is_year_start(self) -> bool: ...
+    @property
+    def is_month_end(self) -> bool: ...
+    @property
+    def is_quarter_end(self) -> bool: ...
+    @property
+    def is_year_end(self) -> bool: ...
+
+    @property
+    def day_of_year(self) -> float: ...
+    @property
+    def dayofyear(self) -> float: ...
+    @property
+    def days_in_month(self) -> float: ...
+    @property
+    def daysinmonth(self) -> float: ...
+    @property
+    def day_of_week(self) -> float: ...
+    @property
+    def dayofweek(self) -> float: ...
+    @property
+    def week(self) -> float: ...
+    @property
+    def weekofyear(self) -> float: ...
+
+    def day_name(self) -> float: ...
+    def month_name(self) -> float: ...
+
+    # error: Return type "float" of "weekday" incompatible with return
+    # type "int" in supertype "date"
+    def weekday(self) -> float: ...  # type: ignore[override]
+
+    # error: Return type "float" of "isoweekday" incompatible with return
+    # type "int" in supertype "date"
+    def isoweekday(self) -> float: ...  # type: ignore[override]
+
+    def total_seconds(self) -> float: ...
+
+    # error: Signature of "today" incompatible with supertype "datetime"
+    def today(self, *args, **kwargs) -> NaTType: ...  # type: ignore[override]
+    # error: Signature of "today" incompatible with supertype "datetime"
+    def now(self, *args, **kwargs) -> NaTType: ...  # type: ignore[override]
+
+    def to_pydatetime(self) -> NaTType: ...
+    def date(self) -> NaTType: ...
+
+    def round(self) -> NaTType: ...
+    def floor(self) -> NaTType: ...
+    def ceil(self) -> NaTType: ...
+
+    def tz_convert(self) -> NaTType: ...
+    def tz_localize(self) -> NaTType: ...
+
+    def replace(self, *args, **kwargs) -> NaTType: ...
+
+    # error: Return type "float" of "year" incompatible with return
+    # type "int" in supertype "date"
+    @property
+    def year(self) -> float: ...  # type: ignore[override]
+
+    @property
+    def quarter(self) -> float: ...
+
+    # error: Return type "float" of "month" incompatible with return
+    # type "int" in supertype "date"
+    @property
+    def month(self) -> float: ...  # type: ignore[override]
+
+    # error: Return type "float" of "day" incompatible with return
+    # type "int" in supertype "date"
+    @property
+    def day(self) -> float: ...  # type: ignore[override]
+
+    # error: Return type "float" of "hour" incompatible with return
+    # type "int" in supertype "date"
+    @property
+    def hour(self) -> float: ...  # type: ignore[override]
+
+    # error: Return type "float" of "minute" incompatible with return
+    # type "int" in supertype "date"
+    @property
+    def minute(self) -> float: ...  # type: ignore[override]
+
+    # error: Return type "float" of "second" incompatible with return
+    # type "int" in supertype "date"
+    @property
+    def second(self) -> float: ...  # type: ignore[override]
+
+    @property
+    def millisecond(self) -> float: ...
+
+    # error: Return type "float" of "microsecond" incompatible with return
+    # type "int" in supertype "date"
+    @property
+    def microsecond(self) -> float: ...  # type: ignore[override]
+
+    @property
+    def nanosecond(self) -> float: ...
+
+    # inject Timedelta properties
+    @property
+    def days(self) -> float: ...
+    @property
+    def microseconds(self) -> float: ...
+    @property
+    def nanoseconds(self) -> float: ...
+
+    # inject Period properties
+    @property
+    def qyear(self) -> float: ...
diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py
@@ -33,6 +33,7 @@
     ArrayLike,
     DtypeObj,
     FrameOrSeriesUnion,
+    Scalar,
 )
 from pandas.util._decorators import doc
 
@@ -763,6 +764,7 @@ def factorize(
         dtype = original.dtype
     else:
         values, dtype = _ensure_data(values)
+        na_value: Scalar
 
         if original.dtype.kind in ["m", "M"]:
             # Note: factorize_array will cast NaT bc it has a __int__
diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py
@@ -38,6 +38,7 @@
 )
 from pandas._libs.tslibs import (
     IncompatibleFrequency,
+    NaTType,
     OutOfBoundsDatetime,
     Timestamp,
     tz_compare,
@@ -2371,7 +2372,7 @@ def __reduce__(self):
     # --------------------------------------------------------------------
     # Null Handling Methods
 
-    _na_value = np.nan
+    _na_value: Union[float, NaTType] = np.nan
     """The expected NA value to use with this index."""
 
     @cache_readonly
diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py
@@ -24,6 +24,7 @@
 )
 from pandas._libs.tslibs import (
     BaseOffset,
+    NaTType,
     Resolution,
     Tick,
 )
@@ -218,7 +219,7 @@ def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs):
 
     _can_hold_na = True
 
-    _na_value = NaT
+    _na_value: NaTType = NaT
     """The expected NA value to use with this index."""
 
     def _convert_tolerance(self, tolerance, target):
diff --git a/pandas/core/internals/construction.py b/pandas/core/internals/construction.py
@@ -368,7 +368,10 @@ def maybe_squeeze_dt64tz(dta: ArrayLike) -> ArrayLike:
     # TODO(EA2D): kludge not needed with 2D EAs
     if isinstance(dta, DatetimeArray) and dta.ndim == 2 and dta.tz is not None:
         assert dta.shape[0] == 1
-        dta = dta[0]
+        # error: Incompatible types in assignment (expression has type
+        # "Union[DatetimeLikeArrayMixin, Union[Any, NaTType]]", variable has
+        # type "Union[ExtensionArray, ndarray]")
+        dta = dta[0]  # type: ignore[assignment]
     return dta
 
 
diff --git a/pandas/core/nanops.py b/pandas/core/nanops.py
@@ -18,6 +18,7 @@
 
 from pandas._libs import (
     NaT,
+    NaTType,
     Timedelta,
     iNaT,
     lib,
@@ -414,11 +415,8 @@ def new_func(
         if datetimelike:
             result = _wrap_results(result, orig_values.dtype, fill_value=iNaT)
             if not skipna:
-                # error: Argument 3 to "_mask_datetimelike_result" has incompatible type
-                # "Optional[ndarray]"; expected "ndarray"
-                result = _mask_datetimelike_result(
-                    result, axis, mask, orig_values  # type: ignore[arg-type]
-                )
+                assert mask is not None  # checked above
+                result = _mask_datetimelike_result(result, axis, mask, orig_values)
 
         return result
 
@@ -601,15 +599,15 @@ def _mask_datetimelike_result(
     axis: Optional[int],
     mask: np.ndarray,
     orig_values: np.ndarray,
-):
+) -> Union[np.ndarray, np.datetime64, np.timedelta64, NaTType]:
     if isinstance(result, np.ndarray):
         # we need to apply the mask
         result = result.astype("i8").view(orig_values.dtype)
         axis_mask = mask.any(axis=axis)
         result[axis_mask] = iNaT
     else:
         if mask.any():
-            result = NaT
+            return NaT
     return result
 
 
@@ -1435,19 +1433,19 @@ def _get_counts(
 
 
 def _maybe_null_out(
-    result: np.ndarray,
+    result: np.ndarray | float | NaTType,
     axis: Optional[int],
     mask: Optional[np.ndarray],
     shape: Tuple[int, ...],
     min_count: int = 1,
-) -> np.ndarray | float:
+) -> np.ndarray | float | NaTType:
     """
     Returns
     -------
     Dtype
         The product of all elements on a given axis. ( NaNs are treated as 1)
     """
-    if mask is not None and axis is not None and getattr(result, "ndim", False):
+    if mask is not None and axis is not None and isinstance(result, np.ndarray):
         null_mask = (mask.shape[axis] - mask.sum(axis) - min_count) < 0
         if np.any(null_mask):
             if is_numeric_dtype(result):
@@ -1461,9 +1459,7 @@ def _maybe_null_out(
                 result[null_mask] = None
     elif result is not NaT:
         if check_below_min_count(shape, mask, min_count):
-            # error: Incompatible types in assignment (expression has type
-            # "float", variable has type "ndarray")
-            result = np.nan  # type: ignore[assignment]
+            result = np.nan
 
     return result
 
diff --git a/pandas/core/tools/datetimes.py b/pandas/core/tools/datetimes.py
@@ -701,7 +701,7 @@ def to_datetime(
     infer_datetime_format: bool = False,
     origin="unix",
     cache: bool = True,
-) -> Union[DatetimeIndex, Series, DatetimeScalar, NaTType]:
+) -> Optional[Union[DatetimeIndex, Series, DatetimeScalar, NaTType]]:
     """
     Convert argument to datetime.
 
diff --git a/pandas/io/excel/_odfreader.py b/pandas/io/excel/_odfreader.py
@@ -202,7 +202,8 @@ def _get_cell_value(self, cell, convert_float: bool) -> Scalar:
         elif cell_type == "time":
             result = pd.to_datetime(str(cell))
             result = cast(pd.Timestamp, result)
-            return result.time()
+            # error: Item "str" of "Union[float, str, NaTType]" has no attribute "time"
+            return result.time()  # type: ignore[union-attr]
         else:
             self.close()
             raise ValueError(f"Unrecognized type {cell_type}")
diff --git a/pandas/io/formats/format.py b/pandas/io/formats/format.py
@@ -1723,7 +1723,10 @@ def _format_datetime64_dateonly(
     if date_format:
         return x.strftime(date_format)
     else:
-        return x._date_repr
+        # error: Item "NaTType" of "Union[NaTType, Any]" has no attribute "_date_repr"
+        #  The underlying problem here is that mypy doesn't understand that NaT
+        #  is a singleton, so that the check above excludes it here.
+        return x._date_repr  # type: ignore[union-attr]
 
 
 def get_format_datetime64(
@@ -1801,13 +1804,15 @@ def get_format_timedelta64(
     consider_values = values_int != iNaT
 
     one_day_nanos = 86400 * 10 ** 9
-    even_days = (
-        # error: Unsupported operand types for % ("ExtensionArray" and "int")
-        np.logical_and(
-            consider_values, values_int % one_day_nanos != 0  # type: ignore[operator]
-        ).sum()
-        == 0
-    )
+    # error: Unsupported operand types for % ("ExtensionArray" and "int")
+    not_midnight = values_int % one_day_nanos != 0  # type: ignore[operator]
+    # error: Argument 1 to "__call__" of "ufunc" has incompatible type
+    # "Union[Any, ExtensionArray, ndarray]"; expected
+    # "Union[Union[int, float, complex, str, bytes, generic],
+    # Sequence[Union[int, float, complex, str, bytes, generic]],
+    # Sequence[Sequence[Any]], _SupportsArray]"
+    both = np.logical_and(consider_values, not_midnight)  # type: ignore[arg-type]
+    even_days = both.sum() == 0
 
     if even_days:
         format = None

Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,7 @@`
`24`	`24`	`)`
`25`	`25`	`from pandas._libs.tslibs import (`
`26`	`26`	`BaseOffset,`
	`27`	`+ NaTType,`
`27`	`28`	`Resolution,`
`28`	`29`	`Tick,`
`29`	`30`	`)`
`@@ -218,7 +219,7 @@ def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs):`
`218`	`219`
`219`	`220`	`_can_hold_na = True`
`220`	`221`
`221`		`- _na_value = NaT`
	`222`	`+ _na_value: NaTType = NaT`
`222`	`223`	`"""The expected NA value to use with this index."""`
`223`	`224`
`224`	`225`	`def _convert_tolerance(self, tolerance, target):`