Move pandas-related types to pandas._typing (#362)

Author: zero323

third_party/3/pyspark/rdd.pyi

@@ -5,7 +5,7 @@ from typing_extensions import Literal
 from numpy import int32, int64, float32, float64, ndarray  # type: ignore
 
 from pyspark._typing import SupportsOrdering
-from pyspark.sql._typing import PandasScalarUDFType, PandasScalarIterUDFType, PandasGroupedMapUDFType, PandasCogroupedMapUDFType, PandasGroupedAggUDFType, PandasMapIterUDFType
+from pyspark.sql.pandas._typing import PandasScalarUDFType, PandasScalarIterUDFType, PandasGroupedMapUDFType, PandasCogroupedMapUDFType, PandasGroupedAggUDFType, PandasMapIterUDFType
 import pyspark.context
 from pyspark.resultiterable import ResultIterable
 from pyspark.serializers import Serializer

third_party/3/pyspark/sql/_typing.pyi

@@ -34,43 +34,6 @@ class SupportsClose(Protocol):
     def close(self, error: Exception) -> None:
         ...
 
-PandasScalarUDFType = Literal[200]
-PandasScalarIterUDFType = Literal[204]
-PandasGroupedMapUDFType = Literal[201]
-PandasCogroupedMapUDFType = Literal[206]
-PandasGroupedAggUDFType = Literal[202]
-PandasMapIterUDFType = Literal[205]
-
-class PandasVariadicScalarToScalarFunction(Protocol):
-    def __call__(self, *_: SeriesLike) -> SeriesLike:
-        ...
-
-PandasScalarToScalarFunction = Union[PandasVariadicScalarToScalarFunction, Callable[[SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike]]
-
-class PandasVariadicScalarToStructFunction(Protocol):
-    def __call__(self, *_: SeriesLike) -> DataFrameLike:
-        ...
-
-PandasScalarToStructFunction = Union[PandasVariadicScalarToStructFunction, Callable[[SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike]]
-
-PandasScalarIterFunction = Callable[[Iterable[Union[SeriesLike, Tuple[SeriesLike, ...], DataFrameLike]]], Iterable[SeriesLike]]
-
-PandasGroupedMapFunction = Union[Callable[[DataFrameLike], DataFrameLike], Callable[[Any, DataFrameLike], DataFrameLike]]
-
-class PandasVariadicGroupedAggFunction(Protocol):
-    def __call__(self, *_: SeriesLike) -> LiteralType:
-        ...
-
-PandasGroupedAggFunction = Union[Callable[[SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], PandasVariadicGroupedAggFunction]
-
-PandasMapIterFunction = Callable[[Iterable[DataFrameLike]], Iterable[DataFrameLike]]
-
-PandasCogroupedMapFunction = Callable[[DataFrameLike, DataFrameLike], DataFrameLike]
-
 class UserDefinedFunctionLike(Protocol):
     def __call__(self, *_: ColumnOrName) -> Column:
         ...
-
-MapIterPandasUserDefinedFunction = NewType("MapIterPandasUserDefinedFunction", FunctionType)
-GroupedMapPandasUserDefinedFunction = NewType("GroupedMapPandasUserDefinedFunction", FunctionType)
-CogroupedMapPandasUserDefinedFunction = NewType("CogroupedMapPandasUserDefinedFunction", FunctionType)

third_party/3/pyspark/sql/dataframe.pyi

@@ -4,7 +4,8 @@ from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, TypeVar
 import pandas.core.frame # type: ignore
 from py4j.java_gateway import JavaObject  # type: ignore
 
-from pyspark.sql._typing import ColumnOrName, LiteralType, MapIterPandasUserDefinedFunction
+from pyspark.sql._typing import ColumnOrName, LiteralType
+from pyspark.sql.pandas._typing import MapIterPandasUserDefinedFunction
 from pyspark.sql.types import *
 from pyspark.sql.context import SQLContext
 from pyspark.sql.group import GroupedData

third_party/3/pyspark/sql/group.pyi

@@ -4,7 +4,8 @@
 from typing import overload
 from typing import Any, Callable, Dict, List, Optional
 
-from pyspark.sql._typing import LiteralType, GroupedMapPandasUserDefinedFunction
+from pyspark.sql._typing import LiteralType
+from pyspark.sql.pandas._typing import GroupedMapPandasUserDefinedFunction
 from pyspark.sql.context import SQLContext
 from pyspark.sql.column import Column
 from pyspark.sql.dataframe import DataFrame

third_party/3/pyspark/sql/pandas/_typing/__init__.pyi

@@ -1,12 +1,52 @@
-# This module contains POC compatibility annotations
+from typing import Any, Callable, Iterable, List, NewType, Optional, Tuple, Type, TypeVar, Union
+from typing_extensions import Protocol, Literal
+from types import FunctionType
+
+from pyspark.sql._typing import LiteralType
+from pyspark.sql.pandas._typing.protocols.frame import DataFrameLike as DataFrameLike
+from pyspark.sql.pandas._typing.protocols.series import SeriesLike as SeriesLike
 
 import pandas.core.frame  # type: ignore[import]
 import pandas.core.series  # type: ignore[import]
 
-from .protocols.frame import DataFrameLike as DataFrameLike
-from .protocols.series import SeriesLike as SeriesLike
-
-from typing import Type
-
+# POC compatibility annotations
 PandasDataFrame: Type[DataFrameLike] = pandas.core.frame.DataFrame
 PandasSeries: Type[SeriesLike] = pandas.core.series.Series
+
+# UDF annotations
+PandasScalarUDFType = Literal[200]
+PandasScalarIterUDFType = Literal[204]
+PandasGroupedMapUDFType = Literal[201]
+PandasCogroupedMapUDFType = Literal[206]
+PandasGroupedAggUDFType = Literal[202]
+PandasMapIterUDFType = Literal[205]
+
+class PandasVariadicScalarToScalarFunction(Protocol):
+    def __call__(self, *_: SeriesLike) -> SeriesLike:
+        ...
+
+PandasScalarToScalarFunction = Union[PandasVariadicScalarToScalarFunction, Callable[[SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], SeriesLike]]
+
+class PandasVariadicScalarToStructFunction(Protocol):
+    def __call__(self, *_: SeriesLike) -> DataFrameLike:
+        ...
+
+PandasScalarToStructFunction = Union[PandasVariadicScalarToStructFunction, Callable[[SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], DataFrameLike]]
+
+PandasScalarIterFunction = Callable[[Iterable[Union[SeriesLike, Tuple[SeriesLike, ...], DataFrameLike]]], Iterable[SeriesLike]]
+
+PandasGroupedMapFunction = Union[Callable[[DataFrameLike], DataFrameLike], Callable[[Any, DataFrameLike], DataFrameLike]]
+
+class PandasVariadicGroupedAggFunction(Protocol):
+    def __call__(self, *_: SeriesLike) -> LiteralType:
+        ...
+
+PandasGroupedAggFunction = Union[Callable[[SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], PandasVariadicGroupedAggFunction]
+
+PandasMapIterFunction = Callable[[Iterable[DataFrameLike]], Iterable[DataFrameLike]]
+
+PandasCogroupedMapFunction = Callable[[DataFrameLike, DataFrameLike], DataFrameLike]
+
+MapIterPandasUserDefinedFunction = NewType("MapIterPandasUserDefinedFunction", FunctionType)
+GroupedMapPandasUserDefinedFunction = NewType("GroupedMapPandasUserDefinedFunction", FunctionType)
+CogroupedMapPandasUserDefinedFunction = NewType("CogroupedMapPandasUserDefinedFunction", FunctionType)

third_party/3/pyspark/sql/pandas/functions.pyi

@@ -1,7 +1,8 @@
 from typing import overload
 from typing import Any, Optional, Union, Callable
 
-from pyspark.sql._typing import AtomicDataTypeOrString, ColumnOrName, DataTypeOrString, GroupedMapPandasUserDefinedFunction, MapIterPandasUserDefinedFunction,CogroupedMapPandasUserDefinedFunction, PandasCogroupedMapFunction, PandasCogroupedMapUDFType, PandasGroupedAggFunction, PandasGroupedAggUDFType, PandasGroupedMapFunction, PandasGroupedMapUDFType, PandasMapIterFunction, PandasMapIterUDFType, PandasScalarIterFunction, PandasScalarIterUDFType, PandasScalarToScalarFunction, PandasScalarToStructFunction, PandasScalarUDFType, UserDefinedFunctionLike
+from pyspark.sql._typing import AtomicDataTypeOrString, ColumnOrName, DataTypeOrString, UserDefinedFunctionLike
+from pyspark.sql.pandas._typing import GroupedMapPandasUserDefinedFunction, MapIterPandasUserDefinedFunction,CogroupedMapPandasUserDefinedFunction, PandasCogroupedMapFunction, PandasCogroupedMapUDFType, PandasGroupedAggFunction, PandasGroupedAggUDFType, PandasGroupedMapFunction, PandasGroupedMapUDFType, PandasMapIterFunction, PandasMapIterUDFType, PandasScalarIterFunction, PandasScalarIterUDFType, PandasScalarToScalarFunction, PandasScalarToStructFunction, PandasScalarUDFType
 
 from pyspark import since as since
 from pyspark.rdd import PythonEvalType as PythonEvalType

third_party/3/pyspark/sql/pandas/group_ops.pyi

@@ -1,6 +1,6 @@
-from typing import Any, Union
+from typing import Union
 
-from pyspark.sql._typing import GroupedMapPandasUserDefinedFunction, CogroupedMapPandasUserDefinedFunction, PandasGroupedMapFunction, PandasCogroupedMapFunction
+from pyspark.sql.pandas._typing import GroupedMapPandasUserDefinedFunction, CogroupedMapPandasUserDefinedFunction, PandasGroupedMapFunction, PandasCogroupedMapFunction
 
 from pyspark import since as since
 from pyspark.rdd import PythonEvalType as PythonEvalType

third_party/3/pyspark/sql/pandas/map_ops.pyi

@@ -1,6 +1,6 @@
-from typing import Any, Union
+from typing import Union
 
-from pyspark.sql._typing import ColumnOrName, LiteralType, MapIterPandasUserDefinedFunction, PandasMapIterFunction
+from pyspark.sql.pandas._typing import  MapIterPandasUserDefinedFunction, PandasMapIterFunction
 from pyspark import since as since
 from pyspark.rdd import PythonEvalType as PythonEvalType
 from pyspark.sql.types import StructType