Repository: modin-project/modin
Pull request: REFACTOR-#4631: Add mypy checks for modin.distributed
Signed-off-by: Rosdyana Kusuma <rosdyana.kusuma@gmail.com>

<!--
Thank you for your contribution!
Please review the contributing docs: https://modin.readthedocs.io/en/latest/development/contributing.html
if you have questions about contributing.
-->

## What do these changes do?

<!-- Please give a short brief about these changes. -->

- [x] commit message follows format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
- [x] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #4631 <!-- issue must be created for each patch -->
- [ ] tests added and passing
- [ ] module layout described at `docs/development/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->


Problem statement:
Add mypy checks for modin.distributed
We need to expand our mypy checks to cover Python files in `modin/distributed/`. There may be a lot of files to modify at once so folks should focus on type annotations one file at a time. Please make sure to add files to `mypy.ini`!

Hints from issue discussion:

Commit message: Fixed mypy issue on partitions.py

Signed-off-by: Rosdyana Kusuma <rosdyana.kusuma@gmail.com>

Patch to review:
diff --git a/modin/distributed/dataframe/pandas/partitions.py b/modin/distributed/dataframe/pandas/partitions.py
index 45c044203a1..2ad5d30cdd9 100644
--- a/modin/distributed/dataframe/pandas/partitions.py
+++ b/modin/distributed/dataframe/pandas/partitions.py
@@ -13,13 +13,14 @@
 
 """Module houses API to operate on Modin DataFrame partitions that are pandas DataFrame(s)."""
 
+from typing import Optional
 import numpy as np
 
 from modin.core.storage_formats.pandas.query_compiler import PandasQueryCompiler
 from modin.pandas.dataframe import DataFrame
 
 
-def unwrap_partitions(api_layer_object, axis=None, get_ip=False):
+def unwrap_partitions(api_layer_object: DataFrame, axis: int = None, get_ip: bool = False) -> list:
     """
     Unwrap partitions of the ``api_layer_object``.
 
@@ -54,10 +55,10 @@ def unwrap_partitions(api_layer_object, axis=None, get_ip=False):
     modin_frame._propagate_index_objs(None)
     if axis is None:
 
-        def _unwrap_partitions():
+        def _unwrap_partitions() -> list:
             [p.drain_call_queue() for p in modin_frame._partitions.flatten()]
 
-            def get_block(partition):
+            def get_block(partition: DataFrame) -> np.ndarray:
                 blocks = partition.list_of_blocks
                 assert (
                     len(blocks) == 1
@@ -99,8 +100,13 @@ def get_block(partition):
 
 
 def from_partitions(
-    partitions, axis, index=None, columns=None, row_lengths=None, column_widths=None
-):
+    partitions: list,
+    axis: int,
+    index: Optional[list] | None = None,
+    columns: Optional[list] | None = None,
+    row_lengths: Optional[list] | None = None,
+    column_widths: Optional[list] | None = None
+) -> DataFrame:
     """
     Create DataFrame from remote partitions.
 
@@ -141,10 +147,16 @@ def from_partitions(
     from modin.core.execution.dispatching.factories.dispatcher import FactoryDispatcher
 
     factory = FactoryDispatcher.get_factory()
-
-    partition_class = factory.io_cls.frame_cls._partition_mgr_cls._partition_class
+    assert factory is not None
+    assert factory.io_cls is not None
+    assert factory.io_cls.frame_cls is not None
+    assert factory.io_cls.frame_cls._partition_mgr_cls is not None
     partition_frame_class = factory.io_cls.frame_cls
+    assert partition_frame_class is not None
     partition_mgr_class = factory.io_cls.frame_cls._partition_mgr_cls
+    assert partition_mgr_class is not None
+    partition_class = factory.io_cls.frame_cls._partition_mgr_cls._partition_class
+    assert partition_class is not None
 
     # Since we store partitions of Modin DataFrame as a 2D NumPy array we need to place
     # passed partitions to 2D NumPy array to pass it to internal Modin Frame class.
diff --git a/mypy.ini b/mypy.ini
index 50f1d90a3ea..09d831bb095 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -21,4 +21,5 @@ warn_unreachable=True
 files = 
         modin/config/*.py,
         modin/logging/*.py,
+        modin/distributed/*.py,
         modin/*.py


Code review comments:
diff_hunk: @@ -99,8 +100,13 @@ def get_block(partition):
 
 
 def from_partitions(
-    partitions, axis, index=None, columns=None, row_lengths=None, column_widths=None
-):
+    partitions: list,
+    axis: int,
path: modin/distributed/dataframe/pandas/partitions.py
text: @user1:
```suggestion
    axis: Optional[int],
```

`axis` can be None.

@author:
Fixed

diff_hunk: @@ -99,8 +100,13 @@ def get_block(partition):
 
 
 def from_partitions(
-    partitions, axis, index=None, columns=None, row_lengths=None, column_widths=None
-):
+    partitions: list,
+    axis: int,
+    index: Optional[list] | None = None,
+    columns: Optional[list] | None = None,
+    row_lengths: Optional[list] | None = None,
+    column_widths: Optional[list] | None = None
path: modin/distributed/dataframe/pandas/partitions.py
text: @user1:
```suggestion
    index: Optional[list] = None,
    columns: Optional[list] = None,
    row_lengths: Optional[list] = None,
    column_widths: Optional[list] = None
```

`Optional` is another way of saying that an object can be `None`.

@author:
Fixed

diff_hunk: @@ -141,10 +147,16 @@ def from_partitions(
     from modin.core.execution.dispatching.factories.dispatcher import FactoryDispatcher
 
     factory = FactoryDispatcher.get_factory()
-
-    partition_class = factory.io_cls.frame_cls._partition_mgr_cls._partition_class
+    assert factory is not None
+    assert factory.io_cls is not None
+    assert factory.io_cls.frame_cls is not None
+    assert factory.io_cls.frame_cls._partition_mgr_cls is not None
path: modin/distributed/dataframe/pandas/partitions.py
text: @user1:
These shouldn't be `None`. Does mypy complain that they could be `None`?

@author:
Hi @user1 this is the complaints from mypy
```
modin/distributed/dataframe/pandas/partitions.py:157:27: error: "None" has no attribute "_partition_mgr_cls"  [attr-defined]

modin/distributed/dataframe/pandas/partitions.py:158:23: error: "None" has no attribute "_partition_mgr_cls"  [attr-defined]
```
Do you think we can ignore it?

@user2:
It may be that we need to add mypy hints to the FactoryDispatcher. 

@user1:
@author could you please add the assertions back so we stop getting these failures, and also add a comment like this:

`# TODO(https://github.com/modin-project/modin/issues/5127): Remove these assertions once the dependencies of this function all have types.`

diff_hunk: @@ -54,10 +55,10 @@ def unwrap_partitions(api_layer_object, axis=None, get_ip=False):
     modin_frame._propagate_index_objs(None)
     if axis is None:
 
-        def _unwrap_partitions():
+        def _unwrap_partitions() -> list:
             [p.drain_call_queue() for p in modin_frame._partitions.flatten()]
 
-            def get_block(partition):
+            def get_block(partition: DataFrame) -> np.ndarray:
path: modin/distributed/dataframe/pandas/partitions.py
text: @user1:
```suggestion
            def get_block(partition: Union[PandasOnRayDataframePartition, PandasOnDaskDataframePartition]) -> np.ndarray:
```

This function is operating on [partitions](https://modin.readthedocs.io/en/stable/development/architecture.html#dataframe-partitioning) of the modin dataframe, i.e. the partitions in `df._query_compiler._modin_frame._partitions`, where `df` is a modin dataframe or series.

@author:
Fixed

diff_hunk: @@ -13,13 +13,14 @@
 
 """Module houses API to operate on Modin DataFrame partitions that are pandas DataFrame(s)."""
 
+from typing import Optional
 import numpy as np
 
 from modin.core.storage_formats.pandas.query_compiler import PandasQueryCompiler
 from modin.pandas.dataframe import DataFrame
 
 
-def unwrap_partitions(api_layer_object, axis=None, get_ip=False):
+def unwrap_partitions(api_layer_object: DataFrame, axis: int = None, get_ip: bool = False) -> list:
path: modin/distributed/dataframe/pandas/partitions.py
text: @user1:
```suggestion
def unwrap_partitions(api_layer_object: Union[DataFrame, Series], axis: int = None, get_ip: bool = False) -> list:
```

This function can take a Series, too.

@author:
Fixed

diff_hunk: @@ -13,13 +13,14 @@
 
 """Module houses API to operate on Modin DataFrame partitions that are pandas DataFrame(s)."""
 
+from typing import Optional
path: modin/distributed/dataframe/pandas/partitions.py
text: @user1:
```suggestion
from typing import Optional, Union
```

diff_hunk: @@ -13,13 +13,14 @@
 
 """Module houses API to operate on Modin DataFrame partitions that are pandas DataFrame(s)."""
 
+from typing import Optional
 import numpy as np
 
 from modin.core.storage_formats.pandas.query_compiler import PandasQueryCompiler
 from modin.pandas.dataframe import DataFrame
path: modin/distributed/dataframe/pandas/partitions.py
text: @user1:
```suggestion
from modin.pandas.dataframe import DataFrame, Series
```