Add minimal mypy type checking (#47)

* Add minimal mypy type checking

.github/workflows/ci.yaml

@@ -4,27 +4,29 @@ on: [push, pull_request]
 
 jobs:
   test:
-
     runs-on: ubuntu-latest
     strategy:
       matrix:
         python-version: [3.6, 3.7, 3.8]
 
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install .[dev]
-    - name: Run pre-commit
-      uses: pre-commit/action@v2.0.0
-    - name: Test with pytest
-      run: |
-        py.test tests -v --cov=rechunker --cov-config .coveragerc --cov-report term-missing
-        coverage xml
-    - name: Codecov
-      uses: codecov/codecov-action@v1
+      - uses: actions/checkout@v2
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install .[dev]
+      - name: Run pre-commit
+        uses: pre-commit/action@v2.0.0
+      - name: Test with pytest
+        run: |
+          py.test tests -v --cov=rechunker --cov-config .coveragerc --cov-report term-missing
+          coverage xml
+      - name: Codecov
+        uses: codecov/codecov-action@v1
+      - name: Check type hints
+        run: |
+          mypy rechunker 

rechunker/_version.py

@@ -51,7 +51,7 @@ class NotThisMethod(Exception):
     """Exception raised if a method is not valid for the current scenario."""
 
 
-LONG_VERSION_PY = {}
+LONG_VERSION_PY = {}  # type: ignore
 HANDLERS = {}
 
 

rechunker/compat.py

@@ -6,7 +6,7 @@
 def prod(iterable: Sequence[int]) -> int:
     """Implementation of `math.prod()` all Python versions."""
     try:
-        from math import prod as mathprod  # Python 3.8
+        from math import prod as mathprod  # type: ignore # Python 3.8
 
         return mathprod(iterable)
     except ImportError:

rechunker/executors/beam.py

@@ -1,5 +1,5 @@
 import uuid
-from typing import Iterable, Mapping, Tuple
+from typing import Iterable, Iterator, Mapping, Tuple
 
 import apache_beam as beam
 
@@ -77,15 +77,15 @@ def expand(self, pcoll):
 
 def _start_stage(
     target_id: str, specs_by_target: Mapping[str, DirectCopySpec],
-) -> Tuple[str, DirectCopySpec]:
+) -> Iterator[Tuple[str, DirectCopySpec]]:
     spec = specs_by_target.get(target_id)
     if spec is not None:
         yield target_id, spec
 
 
 def _copy_tasks(
     target_id: str, spec: DirectCopySpec
-) -> Tuple[str, Tuple[slice, ...], ReadableArray, WriteableArray]:
+) -> Iterator[Tuple[str, Tuple[slice, ...], ReadableArray, WriteableArray]]:
     for key in chunk_keys(spec.source.shape, spec.chunks):
         yield target_id, key, spec.source, spec.target
 

rechunker/executors/python.py

@@ -27,7 +27,7 @@ def prepare_plan(self, specs: Iterable[CopySpec]) -> Task:
                 tasks.append(partial(_direct_array_copy, *direct_spec))
         return partial(_execute_all, tasks)
 
-    def execute_plan(self, plan: Task):
+    def execute_plan(self, plan: Task, **kwargs):
         plan()
 
 

rechunker/executors/pywren.py

@@ -38,7 +38,7 @@ def prepare_plan(self, specs: Iterable[CopySpec]) -> Task:
         # TODO: execute tasks for different specs in parallel
         return partial(_execute_in_series, tasks)
 
-    def execute_plan(self, plan: Task):
+    def execute_plan(self, plan: Task, **kwargs):
         if self.pywren_function_executor is None:
             # No Pywren function executor specified, so use a local one, and shutdown after use
             with pywren_local_function_executor() as pywren_function_executor:

rechunker/executors/util.py

@@ -31,7 +31,7 @@ class DirectCopySpec(NamedTuple):
     chunks: Tuple[int, ...]
 
 
-def split_into_direct_copies(spec: CopySpec) -> Tuple[DirectCopySpec]:
+def split_into_direct_copies(spec: CopySpec) -> Tuple[DirectCopySpec, ...]:
     """Convert a rechunked copy into a list of direct copies."""
     if spec.intermediate.array is None:
         return (DirectCopySpec(spec.read.array, spec.write.array, spec.read.chunks,),)

setup.cfg

@@ -18,6 +18,17 @@ ignore =
     F811,  # redefinition of unused 'loop' from line 10
 max-line-length = 120
 
+[mypy-apache_beam.*]
+ignore_missing_imports = True
+[mypy-dask.*]
+ignore_missing_imports = True
+[mypy-prefect.*]
+ignore_missing_imports = True
+[mypy-pywren_ibm_cloud.*]
+ignore_missing_imports = True
+[mypy-zarr.*]
+ignore_missing_imports = True
+
 # See the docstring in versioneer.py for instructions. Note that you must
 # re-run 'versioneer.py setup' after changing this section, and commit the
 # resulting files.

setup.py

@@ -32,6 +32,7 @@
     "flake8",
     "black",
     "codecov",
+    "mypy==0.782",
 ]
 
 setup(

tests/test_algorithm.py

@@ -64,9 +64,15 @@ def test_consolidate_chunks_limit_error(shape, chunks, itemsize, max_mem, chunk_
         consolidate_chunks(shape, chunks, itemsize, max_mem, chunk_limits=chunk_limits)
 
 
-@pytest.mark.parametrize("shape", [(1000, 50, 1800, 3600),])
-@pytest.mark.parametrize("chunks", [(1, 5, 1800, 3600),])
-@pytest.mark.parametrize("itemsize", [4,])
+@pytest.mark.parametrize(
+    "shape", [(1000, 50, 1800, 3600),],
+)
+@pytest.mark.parametrize(
+    "chunks", [(1, 5, 1800, 3600),],
+)
+@pytest.mark.parametrize(
+    "itemsize", [4,],
+)
 @pytest.mark.parametrize(
     "max_mem, expected",
     [(1_000_000_000, (1, 35, 1800, 3600)), (3_000_000_000, (2, 50, 1800, 3600))],
@@ -183,8 +189,7 @@ def test_rechunking_plan_2d(
 
 @st.composite
 def shapes_chunks_maxmem(draw, ndim=3, itemsize=4, max_len=10_000):
-    """Generate the data we need to test rechunking_plan.
-    """
+    """Generate the data we need to test rechunking_plan."""
     shape = []
     source_chunks = []
     target_chunks = []