ENH: Improve typing of some general functions (#355)

* ENH: Improve typing of some general functions

* Further typing

* ENH: Add overload to factorize

* BUG: Correct npt import

* ENH: Improve unique

* ENH: Imrpveo typing in merge functions, cut and qcut

* CLN: Remove unused imports

* CLN: Catch warning

* TST: Add tests for cut and fixes

* TST: Add final test for cut

* TST: Add tests and fixes for qcut

* TYP: Add typ ignore for overlapping defs

* TST: Add tests for merge

* TST: Add tests for merge_ordered and improve typing

* TST: Add tests for merge_asof and improve typing accuracy

* CLN: Remove TODO since pandas PR opened

* TYP: Final refinements

* TST: Fix intentionally failing test

* MAINT: Refactor MergeHow and add JoinHow

* TYP: Correct use of MergeHow and add test

Co-authored-by: Kevin Sheppard <kevin.sheppard@gmail.com>

Author: bashtage

pandas-stubs/_typing.pyi

@@ -287,7 +287,8 @@ FillnaOptions: TypeAlias = Literal["backfill", "bfill", "ffill", "pad"]
 ReplaceMethod: TypeAlias = Literal["pad", "ffill", "bfill"]
 SortKind: TypeAlias = Literal["quicksort", "mergesort", "heapsort", "stable"]
 NaPosition: TypeAlias = Literal["first", "last"]
-MergeHow: TypeAlias = Literal["left", "right", "outer", "inner"]
+JoinHow: TypeAlias = Literal["left", "right", "outer", "inner"]
+MergeHow: TypeAlias = Union[JoinHow, Literal["cross"]]
 JsonFrameOrient: TypeAlias = Literal[
     "split", "records", "index", "columns", "values", "table"
 ]
@@ -333,4 +334,14 @@ class StyleExportDict(TypedDict, total=False):
 
 CalculationMethod: TypeAlias = Literal["single", "table"]
 
+ValidationOptions: TypeAlias = Literal[
+    "one_to_one",
+    "1:1",
+    "one_to_many",
+    "1:m",
+    "many_to_one",
+    "m:1",
+    "many_to_many",
+    "m:m",
+]
 __all__ = ["npt", "type_t"]

pandas-stubs/core/algorithms.pyi

@@ -1,20 +1,32 @@
 from typing import (
-    Any,
+    Sequence,
     overload,
 )
 
 import numpy as np
+import pandas as pd
 from pandas import (
     Categorical,
+    CategoricalIndex,
     Index,
+    IntervalIndex,
+    PeriodIndex,
     Series,
 )
 from pandas.api.extensions import ExtensionArray
 
 from pandas._typing import AnyArrayLike
 
+# These are type: ignored because the Index types overlap due to inheritance but indices
+# with extension types return the same type while standard type return ndarray
 @overload
-def unique(values: Index) -> Index: ...
+def unique(values: PeriodIndex) -> PeriodIndex: ...  # type: ignore[misc]
+@overload
+def unique(values: CategoricalIndex) -> CategoricalIndex: ...  # type: ignore[misc]
+@overload
+def unique(values: IntervalIndex) -> IntervalIndex: ...  # type: ignore[misc]
+@overload
+def unique(values: Index) -> np.ndarray: ...
 @overload
 def unique(values: Categorical) -> Categorical: ...
 @overload
@@ -23,14 +35,33 @@ def unique(values: Series) -> np.ndarray | ExtensionArray: ...
 def unique(values: np.ndarray | list) -> np.ndarray: ...
 @overload
 def unique(values: ExtensionArray) -> ExtensionArray: ...
+@overload
+def factorize(
+    values: Sequence,
+    sort: bool = ...,
+    # Not actually positional-only, used to handle deprecations in 1.5.0
+    *,
+    use_na_sentinel: bool = ...,
+    size_hint: int | None = ...,
+) -> tuple[np.ndarray, np.ndarray]: ...
+@overload
+def factorize(
+    values: Index | Series,
+    sort: bool = ...,
+    # Not actually positional-only, used to handle deprecations in 1.5.0
+    *,
+    use_na_sentinel: bool = ...,
+    size_hint: int | None = ...,
+) -> tuple[np.ndarray, Index]: ...
+@overload
 def factorize(
-    values: Any,
+    values: Categorical,
     sort: bool = ...,
     # Not actually positional-only, used to handle deprecations in 1.5.0
     *,
     use_na_sentinel: bool = ...,
     size_hint: int | None = ...,
-) -> tuple[np.ndarray, np.ndarray | Index]: ...
+) -> tuple[np.ndarray, Categorical]: ...
 def value_counts(
     values: AnyArrayLike | list | tuple,
     sort: bool = ...,

pandas-stubs/core/frame.pyi

@@ -74,6 +74,7 @@ from pandas._typing import (
     IndexLabel,
     IndexType,
     IntervalClosedType,
+    JoinHow,
     JsonFrameOrient,
     Label,
     Level,
@@ -97,6 +98,7 @@ from pandas._typing import (
     Suffixes,
     T as TType,
     TimestampConvention,
+    ValidationOptions,
     WriteBuffer,
     XMLParsers,
     np_ndarray_bool,
@@ -531,7 +533,7 @@ class DataFrame(NDFrame, OpsMixin):
     def align(
         self,
         other: DataFrame | Series,
-        join: MergeHow = ...,
+        join: JoinHow = ...,
         axis: AxisType | None = ...,
         level: Level | None = ...,
         copy: _bool = ...,
@@ -1101,21 +1103,11 @@ class DataFrame(NDFrame, OpsMixin):
         self,
         other: DataFrame | Series | list[DataFrame | Series],
         on: _str | list[_str] | None = ...,
-        how: MergeHow = ...,
+        how: JoinHow = ...,
         lsuffix: _str = ...,
         rsuffix: _str = ...,
         sort: _bool = ...,
-        validate: Literal[
-            "one_to_one",
-            "1:1",
-            "one_to_many",
-            "1:m",
-            "many_to_one",
-            "m:1",
-            "many_to_many",
-            "m:m",
-        ]
-        | None = ...,
+        validate: ValidationOptions | None = ...,
     ) -> DataFrame: ...
     def merge(
         self,

pandas-stubs/core/reshape/melt.pyi

@@ -1,16 +1,27 @@
+from typing import Hashable
+
 import numpy as np
 from pandas.core.frame import DataFrame
 
+from pandas._typing import HashableT
+
 def melt(
     frame: DataFrame,
     id_vars: tuple | list | np.ndarray | None = ...,
     value_vars: tuple | list | np.ndarray | None = ...,
     var_name: str | None = ...,
-    value_name: str = ...,
+    value_name: Hashable = ...,
     col_level: int | str | None = ...,
     ignore_index: bool = ...,
 ) -> DataFrame: ...
-def lreshape(data: DataFrame, groups, dropna: bool = ..., label=...) -> DataFrame: ...
+def lreshape(
+    data: DataFrame, groups: dict[HashableT, list[HashableT]], dropna: bool = ...
+) -> DataFrame: ...
 def wide_to_long(
-    df: DataFrame, stubnames, i, j, sep: str = ..., suffix: str = ...
+    df: DataFrame,
+    stubnames: str | list[str],
+    i: str | list[str],
+    j: str,
+    sep: str = ...,
+    suffix: str = ...,
 ) -> DataFrame: ...

pandas-stubs/core/reshape/merge.pyi

@@ -1,4 +1,7 @@
-from typing import Sequence
+from typing import (
+    Literal,
+    overload,
+)
 
 from pandas import (
     DataFrame,
@@ -8,131 +11,95 @@ from pandas import (
 from pandas._libs.tslibs import Timedelta
 from pandas._typing import (
     AnyArrayLike,
+    HashableT,
+    JoinHow,
     Label,
+    MergeHow,
+    ValidationOptions,
 )
 
 def merge(
     left: DataFrame | Series,
     right: DataFrame | Series,
-    how: str = ...,
-    on: Label | Sequence | AnyArrayLike | None = ...,
-    left_on: Label | Sequence | AnyArrayLike | None = ...,
-    right_on: Label | Sequence | AnyArrayLike | None = ...,
+    how: MergeHow = ...,
+    on: Label | list[HashableT] | AnyArrayLike | None = ...,
+    left_on: Label | list[HashableT] | AnyArrayLike | None = ...,
+    right_on: Label | list[HashableT] | AnyArrayLike | None = ...,
     left_index: bool = ...,
     right_index: bool = ...,
     sort: bool = ...,
-    suffixes: Sequence[str | None] = ...,
+    suffixes: list[str | None]
+    | tuple[str, str]
+    | tuple[None, str]
+    | tuple[str, None] = ...,
     copy: bool = ...,
     indicator: bool | str = ...,
-    validate: str = ...,
+    validate: ValidationOptions = ...,
 ) -> DataFrame: ...
+@overload
 def merge_ordered(
-    left: DataFrame | Series,
+    left: DataFrame,
+    right: DataFrame,
+    on: Label | list[HashableT] | None = ...,
+    left_on: Label | list[HashableT] | None = ...,
+    right_on: Label | list[HashableT] | None = ...,
+    left_by: Label | list[HashableT] | None = ...,
+    right_by: Label | list[HashableT] | None = ...,
+    fill_method: Literal["ffill"] | None = ...,
+    suffixes: list[str | None]
+    | tuple[str, str]
+    | tuple[None, str]
+    | tuple[str, None] = ...,
+    how: JoinHow = ...,
+) -> DataFrame: ...
+@overload
+def merge_ordered(
+    left: Series,
     right: DataFrame | Series,
-    on: Label | Sequence | AnyArrayLike | None = ...,
-    left_on: Label | Sequence | AnyArrayLike | None = ...,
-    right_on: Label | Sequence | AnyArrayLike | None = ...,
-    left_by: str | list[str] | None = ...,
-    right_by: str | list[str] | None = ...,
-    fill_method: str | None = ...,
-    suffixes: Sequence[str | None] = ...,
-    how: str = ...,
+    on: Label | list[HashableT] | None = ...,
+    left_on: Label | list[HashableT] | None = ...,
+    right_on: Label | list[HashableT] | None = ...,
+    left_by: None = ...,
+    right_by: None = ...,
+    fill_method: Literal["ffill"] | None = ...,
+    suffixes: list[str | None]
+    | tuple[str, str]
+    | tuple[None, str]
+    | tuple[str, None] = ...,
+    how: JoinHow = ...,
+) -> DataFrame: ...
+@overload
+def merge_ordered(
+    left: DataFrame | Series,
+    right: Series,
+    on: Label | list[HashableT] | None = ...,
+    left_on: Label | list[HashableT] | None = ...,
+    right_on: Label | list[HashableT] | None = ...,
+    left_by: None = ...,
+    right_by: None = ...,
+    fill_method: Literal["ffill"] | None = ...,
+    suffixes: list[str | None]
+    | tuple[str, str]
+    | tuple[None, str]
+    | tuple[str, None] = ...,
+    how: JoinHow = ...,
 ) -> DataFrame: ...
 def merge_asof(
     left: DataFrame | Series,
     right: DataFrame | Series,
     on: Label | None = ...,
-    left_on: Label | AnyArrayLike | None = ...,
-    right_on: Label | AnyArrayLike | None = ...,
+    left_on: Label | None = ...,
+    right_on: Label | None = ...,
     left_index: bool = ...,
     right_index: bool = ...,
-    by: str | list[str] | None = ...,
-    left_by: str | None = ...,
-    right_by: str | None = ...,
-    suffixes: Sequence[str | None] = ...,
+    by: Label | list[HashableT] | None = ...,
+    left_by: Label | list[HashableT] | None = ...,
+    right_by: Label | list[HashableT] | None = ...,
+    suffixes: list[str | None]
+    | tuple[str, str]
+    | tuple[None, str]
+    | tuple[str, None] = ...,
     tolerance: int | Timedelta | None = ...,
     allow_exact_matches: bool = ...,
-    direction: str = ...,
+    direction: Literal["backward", "forward", "nearest"] = ...,
 ) -> DataFrame: ...
-
-class _MergeOperation:
-    left = ...
-    right = ...
-    how = ...
-    axis = ...
-    on = ...
-    left_on = ...
-    right_on = ...
-    copy = ...
-    suffixes = ...
-    sort = ...
-    left_index = ...
-    right_index = ...
-    indicator = ...
-    indicator_name = ...
-    def __init__(
-        self,
-        left: Series | DataFrame,
-        right: Series | DataFrame,
-        how: str = ...,
-        on=...,
-        left_on=...,
-        right_on=...,
-        axis=...,
-        left_index: bool = ...,
-        right_index: bool = ...,
-        sort: bool = ...,
-        suffixes=...,
-        copy: bool = ...,
-        indicator: bool = ...,
-        validate=...,
-    ) -> None: ...
-    def get_result(self): ...
-
-class _OrderedMerge(_MergeOperation):
-    fill_method = ...
-    def __init__(
-        self,
-        left,
-        right,
-        on=...,
-        left_on=...,
-        right_on=...,
-        left_index: bool = ...,
-        right_index: bool = ...,
-        axis=...,
-        suffixes=...,
-        copy: bool = ...,
-        fill_method=...,
-        how: str = ...,
-    ) -> None: ...
-    def get_result(self): ...
-
-class _AsOfMerge(_OrderedMerge):
-    by = ...
-    left_by = ...
-    right_by = ...
-    tolerance = ...
-    allow_exact_matches = ...
-    direction = ...
-    def __init__(
-        self,
-        left,
-        right,
-        on=...,
-        left_on=...,
-        right_on=...,
-        left_index: bool = ...,
-        right_index: bool = ...,
-        by=...,
-        left_by=...,
-        right_by=...,
-        axis=...,
-        suffixes=...,
-        copy: bool = ...,
-        fill_method=...,
-        how: str = ...,
-        tolerance=...,
-        allow_exact_matches: bool = ...,
-        direction: str = ...,
-    ) -> None: ...