Fix up some return types for methods with skip_nulls (#215)

Author: MarcoGorelli

spec/API_specification/dataframe_api/_types.py

@@ -23,6 +23,9 @@
 # Type alias: Mypy needs Any, but for readability we need to make clear this
 # is a Python scalar (i.e., an instance of `bool`, `int`, `float`, `str`, etc.)
 Scalar = Any
+# null is a special object which represents a missing value.
+# It is not valid as a type.
+NullType = Any
 
 array = TypeVar("array")
 device = TypeVar("device")

spec/API_specification/dataframe_api/column_object.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
-from typing import Any,NoReturn, Sequence, TYPE_CHECKING, Literal, Generic
+from typing import Any,NoReturn, TYPE_CHECKING, Literal, Generic
 
 from ._types import DType
 
 if TYPE_CHECKING:
-    from . import Bool, null
-    from ._types import Scalar
+    from . import Bool
+    from ._types import NullType, Scalar
 
 
 __all__ = ['Column']
@@ -455,7 +455,7 @@ def __invert__(self: Column[Bool]) -> Column[Bool]:
             If any of the Column's columns is not boolean.
         """
 
-    def any(self: Column[Bool], *, skip_nulls: bool = True) -> bool:
+    def any(self: Column[Bool], *, skip_nulls: bool = True) -> bool | NullType:
         """
         Reduction returns a bool.
 
@@ -465,7 +465,7 @@ def any(self: Column[Bool], *, skip_nulls: bool = True) -> bool:
             If column is not boolean.
         """
 
-    def all(self: Column[Bool], *, skip_nulls: bool = True) -> bool:
+    def all(self: Column[Bool], *, skip_nulls: bool = True) -> bool | NullType:
         """
         Reduction returns a bool.
 
@@ -475,48 +475,48 @@ def all(self: Column[Bool], *, skip_nulls: bool = True) -> bool:
             If column is not boolean.
         """
 
-    def min(self, *, skip_nulls: bool = True) -> Scalar:
+    def min(self, *, skip_nulls: bool = True) -> Scalar | NullType:
         """
         Reduction returns a scalar. Any data type that supports comparisons
         must be supported. The returned value has the same dtype as the column.
         """
 
-    def max(self, *, skip_nulls: bool = True) -> Scalar:
+    def max(self, *, skip_nulls: bool = True) -> Scalar | NullType:
         """
         Reduction returns a scalar. Any data type that supports comparisons
         must be supported. The returned value has the same dtype as the column.
         """
 
-    def sum(self, *, skip_nulls: bool = True) -> Scalar:
+    def sum(self, *, skip_nulls: bool = True) -> Scalar | NullType:
         """
         Reduction returns a scalar. Must be supported for numerical and
         datetime data types. The returned value has the same dtype as the
         column.
         """
 
-    def prod(self, *, skip_nulls: bool = True) -> Scalar:
+    def prod(self, *, skip_nulls: bool = True) -> Scalar | NullType:
         """
         Reduction returns a scalar. Must be supported for numerical data types.
         The returned value has the same dtype as the column.
         """
 
-    def median(self, *, skip_nulls: bool = True) -> Scalar:
+    def median(self, *, skip_nulls: bool = True) -> Scalar | NullType:
         """
         Reduction returns a scalar. Must be supported for numerical and
         datetime data types. Returns a float for numerical data types, and
         datetime (with the appropriate timedelta format string) for datetime
         dtypes.
         """
 
-    def mean(self, *, skip_nulls: bool = True) -> Scalar:
+    def mean(self, *, skip_nulls: bool = True) -> Scalar | NullType:
         """
         Reduction returns a scalar. Must be supported for numerical and
         datetime data types. Returns a float for numerical data types, and
         datetime (with the appropriate timedelta format string) for datetime
         dtypes.
         """
 
-    def std(self, *, correction: int | float = 1, skip_nulls: bool = True) -> Scalar:
+    def std(self, *, correction: int | float = 1, skip_nulls: bool = True) -> Scalar | NullType:
         """
         Reduction returns a scalar. Must be supported for numerical and
         datetime data types. Returns a float for numerical data types, and
@@ -542,7 +542,7 @@ def std(self, *, correction: int | float = 1, skip_nulls: bool = True) -> Scalar
             Whether to skip null values.
         """
 
-    def var(self, *, correction: int | float = 1, skip_nulls: bool = True) -> Scalar:
+    def var(self, *, correction: int | float = 1, skip_nulls: bool = True) -> Scalar | NullType:
         """
         Reduction returns a scalar. Must be supported for numerical and
         datetime data types. Returns a float for numerical data types, and
@@ -662,7 +662,7 @@ def unique_indices(self, *, skip_nulls: bool = True) -> Column[Any]:
         """
         ...
 
-    def fill_nan(self: Column[DType], value: float | 'null', /) -> Column[DType]:
+    def fill_nan(self: Column[DType], value: float | NullType, /) -> Column[DType]:
         """
         Fill floating point ``nan`` values with the given fill value.
 

spec/API_specification/dataframe_api/dataframe_object.py

@@ -6,8 +6,8 @@
 if TYPE_CHECKING:
     from .column_object import Column
     from .groupby_object import GroupBy
-    from . import Bool, null
-    from ._types import Scalar
+    from . import Bool
+    from ._types import NullType, Scalar
 
 
 __all__ = ["DataFrame"]
@@ -768,7 +768,7 @@ def unique_indices(self, keys: Sequence[str], *, skip_nulls: bool = True) -> Col
         """
         ...
 
-    def fill_nan(self, value: float | 'null', /) -> DataFrame:
+    def fill_nan(self, value: float | NullType, /) -> DataFrame:
         """
         Fill ``nan`` values with the given fill value.
 

spec/conf.py

@@ -84,6 +84,7 @@
     ('py:class', 'Scalar'),
     ('py:class', 'Bool'),
     ('py:class', 'optional'),
+    ('py:class', 'NullType'),
 ]
 # NOTE: this alias handling isn't used yet - added in anticipation of future
 #       need based on dataframe API aliases.

spec/design_topics/python_builtin_types.md

@@ -18,7 +18,7 @@ class DataFrame:
         ...
 
 class Column:
-    def mean(self, skip_nulls: bool = True) -> float:
+    def mean(self, skip_nulls: bool = True) -> float | NullType:
         ...
 
 larger = df2 > df1.get_column_by_name('foo').mean()