Fix two-argument math functions

zero323 · Jul 9, 2019 · eeb854b · eeb854b
1 parent eee964a
commit eeb854b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/third_party/3/pyspark/sql/_typing.pyi b/third_party/3/pyspark/sql/_typing.pyi
@@ -1,4 +1,4 @@
-from typing import Any, List, Optional, TypeVar, Union
+from typing import Any, List, Optional, TypeVar, Union, SupportsFloat
 from typing_extensions import Protocol
 import datetime
 import decimal
@@ -7,6 +7,7 @@ import pyspark.sql.column
 import pyspark.sql.types
 
 ColumnOrName = Union[pyspark.sql.column.Column, str]
+ColumnNameOrFloat = Union[ColumnOrName, SupportsFloat]
 DecimalLiteral = decimal.Decimal
 DateTimeLiteral = Union[datetime.datetime, datetime.date]
 LiteralType = Union[bool, int, float, str]

diff --git a/third_party/3/pyspark/sql/functions.pyi b/third_party/3/pyspark/sql/functions.pyi
@@ -6,7 +6,7 @@ from typing import Any, Optional, Union, Dict, Callable
 import pandas.core.frame   # type: ignore
 import pandas.core.series     # type: ignore
 
-from pyspark.sql._typing import ColumnOrName
+from pyspark.sql._typing import ColumnOrName, ColumnNameOrFloat
 from pyspark.sql.column import Column
 from pyspark.sql.dataframe import DataFrame
 from pyspark.sql.types import ArrayType, DataType, StructField, StructType
@@ -159,7 +159,7 @@ def asc_nulls_last(col: ColumnOrName) -> Column: ...
 def ascii(col: ColumnOrName) -> Column: ...
 def asin(col: ColumnOrName) -> Column: ...
 def atan(col: ColumnOrName) -> Column: ...
-def atan2(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
+def atan2(col1: ColumnNameOrFloat, col2: ColumnNameOrFloat) -> Column: ...
 def avg(col: ColumnOrName) -> Column: ...
 def base64(col: ColumnOrName) -> Column: ...
 def bitwiseNOT(col: ColumnOrName) -> Column: ...
@@ -181,7 +181,7 @@ def desc_nulls_last(col: ColumnOrName) -> Column: ...
 def exp(col: ColumnOrName) -> Column: ...
 def expm1(col: ColumnOrName) -> Column: ...
 def floor(col: ColumnOrName) -> Column: ...
-def hypot(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
+def hypot(col1: ColumnNameOrFloat, col2: ColumnNameOrFloat) -> Column: ...
 def kurtosis(col: ColumnOrName) -> Column: ...
 def lit(col: Any) -> Column: ...
 def log10(col: ColumnOrName) -> Column: ...
@@ -192,7 +192,7 @@ def max(col: ColumnOrName) -> Column: ...
 def mean(col: ColumnOrName) -> Column: ...
 def min(col: ColumnOrName) -> Column: ...
 def percent_rank() -> Column: ...
-def pow(col1: ColumnOrName, col2: Union[ColumnOrName, float]) -> Column: ...
+def pow(col1: ColumnNameOrFloat, col2: ColumnNameOrFloat) -> Column: ...
 def radians(col: ColumnOrName) -> Column: ...
 def rank() -> Column: ...
 def rint(col: ColumnOrName) -> Column: ...