Repository: modin-project/modin

Issue:
REFACTOR: Add mypy checks for modin.core.dataframe.base



Discussion:
Could you please assign me this issue?
@shalearkane sure. I'm happy to help you merge a PR for this issue. You can follow the instructions [here](https://github.com/modin-project/modin/issues/5073#issue-1395324747), except replace `modin/core/dataframe/algebra` with `modin/core/dataframe/base`.

Resolving patch:
diff --git a/modin/core/dataframe/base/interchange/dataframe_protocol/dataframe.py b/modin/core/dataframe/base/interchange/dataframe_protocol/dataframe.py
--- a/modin/core/dataframe/base/interchange/dataframe_protocol/dataframe.py
+++ b/modin/core/dataframe/base/interchange/dataframe_protocol/dataframe.py
@@ -18,9 +18,9 @@
 """
 
 from abc import ABC, abstractmethod
-from typing import Optional, Iterable, Sequence, Any, Tuple, Dict, TypedDict
+from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, TypedDict
 
-from .utils import DlpackDeviceType, DTypeKind, ColumnNullType
+from .utils import ColumnNullType, DlpackDeviceType, DTypeKind
 
 
 class ColumnBuffers(TypedDict):  # noqa: GL08
@@ -89,7 +89,7 @@ def ptr(self) -> int:
         pass
 
     @abstractmethod
-    def __dlpack__(self):
+    def __dlpack__(self) -> Any:
         """
         Produce DLPack capsule (see array API standard).
 
diff --git a/modin/core/dataframe/base/interchange/dataframe_protocol/utils.py b/modin/core/dataframe/base/interchange/dataframe_protocol/utils.py
--- a/modin/core/dataframe/base/interchange/dataframe_protocol/utils.py
+++ b/modin/core/dataframe/base/interchange/dataframe_protocol/utils.py
@@ -19,9 +19,10 @@
 
 import enum
 import re
-import pandas
-import numpy as np
+from typing import Optional, Union
 
+import numpy as np
+import pandas
 from pandas.api.types import is_datetime64_dtype
 
 
@@ -134,7 +135,7 @@ class Endianness:
     NA = "|"
 
 
-def pandas_dtype_to_arrow_c(dtype) -> str:
+def pandas_dtype_to_arrow_c(dtype: Union[np.dtype, pandas.CategoricalDtype]) -> str:
     """
     Represent pandas `dtype` as a format string in Apache Arrow C notation.
 
@@ -168,7 +169,7 @@ def pandas_dtype_to_arrow_c(dtype) -> str:
     )
 
 
-def raise_copy_alert(copy_reason=None):
+def raise_copy_alert(copy_reason: Optional[str] = None) -> None:
     """
     Raise a ``RuntimeError`` mentioning that there's a copy required.
 
diff --git a/modin/core/dataframe/base/partitioning/axis_partition.py b/modin/core/dataframe/base/partitioning/axis_partition.py
--- a/modin/core/dataframe/base/partitioning/axis_partition.py
+++ b/modin/core/dataframe/base/partitioning/axis_partition.py
@@ -13,7 +13,8 @@
 
 """Base class of an axis partition for a Modin Dataframe."""
 
-from abc import ABC
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Iterable, Optional, Tuple, Union
 
 
 class BaseDataframeAxisPartition(ABC):  # pragma: no cover
@@ -23,17 +24,23 @@ class BaseDataframeAxisPartition(ABC):  # pragma: no cover
     This class is intended to simplify the way that operations are performed.
     """
 
+    @property
+    @abstractmethod
+    def list_of_blocks(self) -> list:
+        """Get the list of physical partition objects that compose this partition."""
+        pass
+
     def apply(
         self,
-        func,
-        *args,
-        num_splits=None,
-        other_axis_partition=None,
-        maintain_partitioning=True,
-        lengths=None,
-        manual_partition=False,
-        **kwargs,
-    ):
+        func: Callable,
+        *args: Iterable,
+        num_splits: Optional[int] = None,
+        other_axis_partition: Optional["BaseDataframeAxisPartition"] = None,
+        maintain_partitioning: bool = True,
+        lengths: Optional[Iterable] = None,
+        manual_partition: bool = False,
+        **kwargs: dict,
+    ) -> Any:
         """
         Apply a function to this axis partition along full axis.
 
@@ -81,7 +88,7 @@ def apply(
     instance_type = None
     partition_type = None
 
-    def _wrap_partitions(self, partitions):
+    def _wrap_partitions(self, partitions: list) -> list:
         """
         Wrap remote partition objects with `BaseDataframePartition` class.
 
@@ -95,9 +102,14 @@ def _wrap_partitions(self, partitions):
         list
             List of wrapped remote partition objects.
         """
+        assert self.partition_type is not None
+        assert self.instance_type is not None  # type: ignore
+
         return [self.partition_type(obj) for obj in partitions]
 
-    def force_materialization(self, get_ip=False):
+    def force_materialization(
+        self, get_ip: bool = False
+    ) -> "BaseDataframeAxisPartition":
         """
         Materialize axis partitions into a single partition.
 
@@ -114,9 +126,11 @@ def force_materialization(self, get_ip=False):
         materialized = self.apply(
             lambda x: x, num_splits=1, maintain_partitioning=False
         )
-        return type(self)(materialized, get_ip=get_ip)
+        return type(self)(materialized, get_ip=get_ip)  # type: ignore[call-arg]
 
-    def unwrap(self, squeeze=False, get_ip=False):
+    def unwrap(
+        self, squeeze: bool = False, get_ip: bool = False
+    ) -> Union[list, Tuple[list, list]]:
         """
         Unwrap partitions from this axis partition.
 
@@ -134,17 +148,19 @@ def unwrap(self, squeeze=False, get_ip=False):
 
         Notes
         -----
-        If `get_ip=True`, a list of tuples of Ray.ObjectRef/Dask.Future to node ip addresses and
+        If `get_ip=True`, a tuple of lists of Ray.ObjectRef/Dask.Future to node ip addresses and
         unwrapped partitions, respectively, is returned if Ray/Dask is used as an engine
         (i.e. [(Ray.ObjectRef/Dask.Future, Ray.ObjectRef/Dask.Future), ...]).
         """
         if squeeze and len(self.list_of_blocks) == 1:
             if get_ip:
-                return self.list_of_ips[0], self.list_of_blocks[0]
+                # TODO(https://github.com/modin-project/modin/issues/5176): Stop ignoring the list_of_ips
+                # check once we know that we're not calling list_of_ips on python axis partitions
+                return self.list_of_ips[0], self.list_of_blocks[0]  # type: ignore[attr-defined]
             else:
                 return self.list_of_blocks[0]
         else:
             if get_ip:
-                return list(zip(self.list_of_ips, self.list_of_blocks))
+                return list(zip(self.list_of_ips, self.list_of_blocks))  # type: ignore[attr-defined]
             else:
                 return self.list_of_blocks


Test patch: