Add static type checking via Mypy (#6381)

Adds static type checking to cuDF Python via MyPy.

* An additional `mypy` style check is enabled in CI
* `mypy` is run as part of the pre-commit hook
* Many parts of the cuDF internal code now have type annotations
* Any new internal code is expected to be written with type annotations (not public-facing APIs)

Authors:
  - Ashwin Srinath (@shwina)

Approvers:
  - Dillon Cullinan (@dillon-cullinan)
  - Keith Kraus (@kkraus14)
  - Christopher Harris (@cwharris)

URL: #6381

Author: shwina

.pre-commit-config.yaml

@@ -32,6 +32,15 @@ repos:
                 language: system
                 files: \.(cu|cuh|h|hpp|cpp|inl)$
                 args: ['-fallback-style=none']
+      - repo: local
+        hooks:
+              - id: mypy
+                name: mypy
+                description: mypy
+                pass_filenames: false
+                entry: mypy --config-file=python/cudf/setup.cfg python/cudf/cudf
+                language: system
+                types: [python]
 
 default_language_version:
       python: python3

ci/checks/style.sh

@@ -29,6 +29,10 @@ FLAKE_RETVAL=$?
 FLAKE_CYTHON=`flake8 --config=python/.flake8.cython`
 FLAKE_CYTHON_RETVAL=$?
 
+# Run mypy and get results/return code
+MYPY_CUDF=`mypy --config=python/cudf/setup.cfg python/cudf/cudf`
+MYPY_CUDF_RETVAL=$?
+
 # Run clang-format and check for a consistent code format
 CLANG_FORMAT=`python cpp/scripts/run-clang-format.py 2>&1`
 CLANG_FORMAT_RETVAL=$?
@@ -66,6 +70,14 @@ else
   echo -e "\n\n>>>> PASSED: flake8-cython style check\n\n"
 fi
 
+if [ "$MYPY_CUDF_RETVAL" != "0" ]; then
+  echo -e "\n\n>>>> FAILED: mypy style check; begin output\n\n"
+  echo -e "$MYPY_CUDF"
+  echo -e "\n\n>>>> FAILED: mypy style check; end output\n\n"
+else
+  echo -e "\n\n>>>> PASSED: mypy style check\n\n"
+fi
+
 if [ "$CLANG_FORMAT_RETVAL" != "0" ]; then
   echo -e "\n\n>>>> FAILED: clang format check; begin output\n\n"
   echo -e "$CLANG_FORMAT"
@@ -79,7 +91,7 @@ HEADER_META=`ci/checks/headers_test.sh`
 HEADER_META_RETVAL=$?
 echo -e "$HEADER_META"
 
-RETVALS=($ISORT_RETVAL $BLACK_RETVAL $FLAKE_RETVAL $FLAKE_CYTHON_RETVAL $CLANG_FORMAT_RETVAL $HEADER_META_RETVAL)
+RETVALS=($ISORT_RETVAL $BLACK_RETVAL $FLAKE_RETVAL $FLAKE_CYTHON_RETVAL $CLANG_FORMAT_RETVAL $HEADER_META_RETVAL $MYPY_CUDF_RETVAL)
 IFS=$'\n'
 RETVAL=`echo "${RETVALS[*]}" | sort -nr | head -n1`
 

conda/environments/cudf_dev_cuda10.1.yml

@@ -40,6 +40,8 @@ dependencies:
   - flake8=3.8.3
   - black=19.10
   - isort=5.0.7
+  - mypy=0.782
+  - typing_extensions
   - pre_commit
   - dask>=2.22.0
   - distributed>=2.22.0

conda/environments/cudf_dev_cuda10.2.yml

@@ -40,6 +40,8 @@ dependencies:
   - flake8=3.8.3
   - black=19.10
   - isort=5.0.7
+  - mypy=0.782
+  - typing_extensions
   - pre_commit
   - dask>=2.22.0
   - distributed>=2.22.0

conda/environments/cudf_dev_cuda11.0.yml

@@ -40,6 +40,8 @@ dependencies:
   - flake8=3.8.3
   - black=19.10
   - isort=5.0.7
+  - mypy=0.782
+  - typing_extensions
   - pre_commit
   - dask>=2.22.0
   - distributed>=2.22.0

conda/recipes/cudf/meta.yaml

@@ -34,6 +34,7 @@ requirements:
   run:
     - protobuf
     - python
+    - typing_extensions
     - pandas >=1.0,<1.2.0dev0
     - cupy >7.1.0,<9.0.0a0
     - numba >=0.49.0

python/cudf/cudf/_lib/__init__.py

@@ -10,13 +10,16 @@
     datetime,
     filling,
     gpuarrow,
+    groupby,
     hash,
     interop,
     join,
+    json,
     merge,
     null_mask,
     nvtext,
     orc,
+    parquet,
     partitioning,
     quantiles,
     reduce,
@@ -27,6 +30,7 @@
     search,
     sort,
     stream_compaction,
+    string_casting,
     strings,
     table,
     transpose,

python/cudf/cudf/_lib/column.pyi

@@ -0,0 +1,124 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+
+from __future__ import annotations
+from typing import Tuple, Union, TypeVar, Optional
+
+from cudf._typing import DtypeObj, Dtype, ScalarLike
+from cudf.core.buffer import Buffer
+from cudf.core.column import ColumnBase
+
+
+T = TypeVar("T")
+
+class Column:
+    _data: Optional[Buffer]
+    _mask: Optional[Buffer]
+    _base_data: Optional[Buffer]
+    _base_mask: Optional[Buffer]
+    _dtype: DtypeObj
+    _offset: int
+    _null_count: int
+    _children: Tuple[ColumnBase, ...]
+    _base_children: Tuple[ColumnBase, ...]
+
+    def __init__(
+        self,
+        data: Optional[Buffer],
+        dtype: Dtype,
+        size: int = None,
+        mask: Optional[Buffer] = None,
+        offset: int = None,
+        null_count: int = None,
+        children: Tuple[ColumnBase, ...] = (),
+    ) -> None:
+        ...
+
+    @property
+    def base_size(self) -> int:
+        ...
+
+    @property
+    def dtype(self) -> DtypeObj:
+        ...
+
+    @property
+    def size(self) -> int:
+        ...
+
+    @property
+    def base_data(self) -> Optional[Buffer]:
+        ...
+
+    @property
+    def base_data_ptr(self) -> int:
+        ...
+
+    @property
+    def data(self) -> Optional[Buffer]:
+        ...
+
+    @property
+    def data_ptr(self) -> int:
+        ...
+
+    def set_base_data(self, value: Buffer) -> None:
+        ...
+
+    @property
+    def nullable(self) -> bool:
+        ...
+
+    @property
+    def has_nulls(self) -> bool:
+        ...
+
+    @property
+    def base_mask(self) -> Optional[Buffer]:
+        ...
+
+    @property
+    def base_mask_ptr(self) -> int:
+        ...
+
+    @property
+    def mask(self) -> Optional[Buffer]:
+        ...
+
+    @property
+    def mask_ptr(self) -> int:
+        ...
+
+    def set_base_mask(self, value: Optional[Buffer]) -> None:
+        ...
+
+    def set_mask(self: T, value: Optional[Buffer]) -> T:
+        ...
+
+    @property
+    def null_count(self) -> int:
+        ...
+
+    @property
+    def offset(self) -> int:
+        ...
+
+    @property
+    def base_children(self) -> Tuple[ColumnBase, ...]:
+        ...
+
+    @property
+    def children(self) -> Tuple[ColumnBase, ...]:
+        ...
+
+    def set_base_children(self, value: Tuple[ColumnBase, ...]) -> None:
+        ...
+
+    def _mimic_inplace(self, other_col: ColumnBase, inplace=False) -> Optional[ColumnBase]:
+        ...
+
+    @staticmethod
+    def from_scalar(
+        val: ScalarLike,
+        size: int
+    ) -> ColumnBase:  # TODO: This should be Scalar, not ScalarLike
+        ...

python/cudf/cudf/_lib/column.pyx

@@ -60,14 +60,14 @@ cdef class Column:
     The *dtype* indicates the Column's element type.
     """
     def __init__(
-            self,
-            object data,
-            int size,
-            object dtype,
-            object mask=None,
-            int offset=0,
-            object null_count=None,
-            object children=()
+        self,
+        object data,
+        int size,
+        object dtype,
+        object mask=None,
+        int offset=0,
+        object null_count=None,
+        object children=()
     ):
 
         self._size = size
@@ -247,10 +247,10 @@ cdef class Column:
             )
 
         return cudf.core.column.build_column(
-            self.data,
-            self.dtype,
-            mask,
-            self.size,
+            data=self.data,
+            dtype=self.dtype,
+            mask=mask,
+            size=self.size,
             offset=0,
             children=self.children
         )
@@ -561,25 +561,22 @@ cdef class Column:
         children = tuple(children)
 
         result = cudf.core.column.build_column(
-            data,
-            dtype,
-            mask,
-            size,
-            offset,
-            null_count,
-            tuple(children)
+            data=data,
+            dtype=dtype,
+            mask=mask,
+            size=size,
+            offset=offset,
+            null_count=null_count,
+            children=tuple(children)
         )
 
         return result
 
-
-def make_column_from_scalar(object py_val, size_type size):
-
-    cdef DeviceScalar val = py_val.device_value
-
-    cdef const scalar* c_val = val.get_raw_ptr()
-    cdef unique_ptr[column] c_result
-    with nogil:
-        c_result = move(cpp_make_column_from_scalar(c_val[0], size))
-
-    return Column.from_unique_ptr(move(c_result))
+    @staticmethod
+    def from_scalar(py_val, size_type size):
+        cdef DeviceScalar val = py_val.device_value
+        cdef const scalar* c_val = val.get_raw_ptr()
+        cdef unique_ptr[column] c_result
+        with nogil:
+            c_result = move(cpp_make_column_from_scalar(c_val[0], size))
+        return Column.from_unique_ptr(move(c_result))

python/cudf/cudf/_lib/table.pyi

@@ -0,0 +1,29 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+
+from typing import List, Any, Optional, TYPE_CHECKING
+
+import cudf
+
+class Table(object):
+    _data: cudf.core.column_accessor.ColumnAccessor
+    _index: Optional[cudf.core.index.Index]
+
+    def __init__(self, data: object = None, index: object = None) -> None: ...
+
+    @property
+    def _num_columns(self) -> int: ...
+
+    @property
+    def _num_indices(self) -> int: ...
+
+    @property
+    def _num_rows(self) -> int: ...
+
+    @property
+    def _column_names(self) -> List[Any]: ...
+
+    @property
+    def _index_names(self) -> List[Any]: ...
+
+    @property
+    def _columns(self) -> List[Any]: ... # TODO: actually, a list of columns

python/cudf/cudf/_typing.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+
+from typing import TYPE_CHECKING, Any, TypeVar, Union
+
+import numpy as np
+from pandas import Period, Timedelta, Timestamp
+from pandas.api.extensions import ExtensionDtype
+
+if TYPE_CHECKING:
+    import cudf
+
+# Many of these are from
+# https://github.com/pandas-dev/pandas/blob/master/pandas/_typing.py
+
+Dtype = Union["ExtensionDtype", str, np.dtype]
+DtypeObj = Union["ExtensionDtype", np.dtype]
+
+# scalars
+DatetimeLikeScalar = TypeVar(
+    "DatetimeLikeScalar", Period, Timestamp, Timedelta
+)
+ScalarLike = Any
+
+# columns
+ColumnLike = Any
+
+# binary operation
+BinaryOperand = Union["cudf.Scalar", "cudf.core.column.ColumnBase"]

python/cudf/cudf/core/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) 2018-2020, NVIDIA CORPORATION.
 
-from cudf.core import buffer, column, common
+from cudf.core import buffer, column, column_accessor, common
 from cudf.core.buffer import Buffer
 from cudf.core.dataframe import DataFrame, from_pandas, merge
 from cudf.core.index import (

python/cudf/cudf/core/abc.py

@@ -12,9 +12,9 @@
     try:
         import pickle5 as pickle
     except ImportError:
-        import pickle
+        import pickle  # type: ignore
 else:
-    import pickle
+    import pickle  # type: ignore
 
 
 class Serializable(abc.ABC):