HyukjinKwon commented on a change in pull request #33954:
URL: https://github.com/apache/spark/pull/33954#discussion_r707829443
##########
File path: python/pyspark/pandas/typedef/typehints.py
##########
@@ -538,6 +634,165 @@ def infer_return_type(f: Callable) -> Union[SeriesType,
DataFrameType, ScalarTyp
return ScalarType(*types)
+# TODO: once pandas exposes a typing module like numpy.typing, we should
deprecate
+# this logic and migrate to it with implementing the typing module in pandas
API on Spark.
+
+
+def create_type_for_series_type(param: Any) -> Type[SeriesType]:
+ """
+ Supported syntax:
+
+ >>> str(pd.Series[float]).endswith("SeriesType[float]")
+ True
+ """
+ from pyspark.pandas.typedef import NameTypeHolder
+
+ if isinstance(param, ExtensionDtype):
+ new_class = type("NameType", (NameTypeHolder,), {}) # type:
Type[NameTypeHolder]
+ new_class.tpe = param
+ else:
+ new_class = param.type if isinstance(param, np.dtype) else param
+
+ return SeriesType[new_class] # type: ignore
+
+
+# TODO: Remove this variadic-generic hack by tuple once ww drop Python up to
3.9.
+# See also PEP 646. One problem is that pandas doesn't inherits Generic[T]
+# so we might have to leave this hack only for monkey-patching pandas
DataFrame.
+def create_tuple_for_frame_type(params: Any) -> object:
+ """
+ This is a workaround to support variadic generic in DataFrame.
+
+ See https://github.com/python/typing/issues/193
+ we always wraps the given type hints by a tuple to mimic the variadic
generic.
+
+ Supported syntax:
+
+ >>> import pandas as pd
+ >>> pdf = pd.DataFrame({'a': range(1)})
+
+ Typing data columns only:
+
+ >>> pd.DataFrame[float, float]
Review comment:
maybe I should use ps to make it less confusing in the example.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
