Validate type hints with mypy, and use mypy in our CI (#164)

* Add mypy support to s3torchconnector

* Add mypy type checking to python-checks.yml

* Update DEVELOPMENT.md to add mypy tests

---------

Co-authored-by: Simon Beal <[email protected]>

Author: muddyfish

.github/workflows/python-checks.yml

@@ -88,18 +88,25 @@ jobs:
       - name: Install Python dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install flake8 black
+          python -m pip install flake8 black mypy
+          python -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
+          python -m pip install ./s3torchconnectorclient
+          python -m pip install ./s3torchconnector
+
       - name: Lint with flake8
         run: |
           # stop the build if there are Python syntax errors or undefined names
           flake8 s3torchconnector/ --count --select=E9,F63,F7,F82 --show-source --statistics
           flake8 s3torchconnectorclient/python --count --select=E9,F63,F7,F82 --show-source --statistics
-
       - name: Lint with Black
         uses: psf/black@stable
         with:
           options: "--check --verbose"
           src: "."
+      - name: Typecheck with mypy
+        run: |
+          mypy s3torchconnector/src
+          mypy s3torchconnectorclient/python/src
 
   dependencies:
     name: Python dependencies checks

doc/DEVELOPMENT.md

@@ -67,9 +67,11 @@ pushing new Rust commits.
 
 For Python code changes, run 
 ```bash
-black --verbose python/
+black --verbose .
 flake8 s3torchconnector/ --count --select=E9,F63,F7,F82 --show-source --statistics
 flake8 s3torchconnectorclient/python --count --select=E9,F63,F7,F82 --show-source --statistics
+mypy s3torchconnector/src
+mypy s3torchconnectorclient/python/src
 ```
  to lint.
 

s3torchconnector/pyproject.toml

@@ -32,7 +32,8 @@ test = [
     "pytest-timeout",
     "hypothesis",
     "flake8",
-    "black"
+    "black",
+    "mypy"
 ]
 
 e2e = [

s3torchconnector/src/s3torchconnector/_s3client/_mock_s3client.py

@@ -1,6 +1,8 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
 
+from typing import Optional
+
 from s3torchconnectorclient._mountpoint_s3_client import (
     MockMountpointS3Client,
     MountpointS3Client,
@@ -21,7 +23,7 @@ def __init__(
         region: str,
         bucket: str,
         part_size: int = 8 * 1024 * 1024,
-        user_agent: UserAgent = None,
+        user_agent: Optional[UserAgent] = None,
     ):
         super().__init__(region, user_agent=user_agent)
         self._mock_client = MockMountpointS3Client(

s3torchconnector/src/s3torchconnector/_s3client/_s3client.py

@@ -32,11 +32,16 @@ def _identity(obj: Any) -> Any:
 
 
 class S3Client:
-    def __init__(self, region: str, endpoint: str = None, user_agent: UserAgent = None):
+    def __init__(
+        self,
+        region: str,
+        endpoint: Optional[str] = None,
+        user_agent: Optional[UserAgent] = None,
+    ):
         self._region = region
         self._endpoint = endpoint
-        self._real_client = None
-        self._client_pid = None
+        self._real_client: Optional[MountpointS3Client] = None
+        self._client_pid: Optional[int] = None
         user_agent = user_agent or UserAgent()
         self._user_agent_prefix = user_agent.prefix
 
@@ -46,6 +51,7 @@ def _client(self) -> MountpointS3Client:
             self._client_pid = os.getpid()
             # `MountpointS3Client` does not survive forking, so re-create it if the PID has changed.
             self._real_client = self._client_builder()
+        assert self._real_client is not None
         return self._real_client
 
     @property

s3torchconnector/src/s3torchconnector/_user_agent.py

@@ -1,14 +1,14 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
-from typing import List
+from typing import List, Optional
 
 from ._version import __version__
 
 # https://www.rfc-editor.org/rfc/rfc9110#name-user-agent
 
 
 class UserAgent:
-    def __init__(self, comments: List[str] = None):
+    def __init__(self, comments: Optional[List[str]] = None):
         if comments is not None and not isinstance(comments, list):
             raise ValueError("Argument comments must be a List[str]")
         self._user_agent_prefix = f"{__package__}/{__version__}"

s3torchconnector/src/s3torchconnector/s3checkpoint.py

@@ -1,5 +1,6 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
+from typing import Optional
 
 from ._s3dataset_common import parse_s3_uri
 from ._s3client import S3Client
@@ -16,7 +17,7 @@ class S3Checkpoint:
     torch.load, and torch.save.
     """
 
-    def __init__(self, region: str, endpoint: str = None):
+    def __init__(self, region: str, endpoint: Optional[str] = None):
         self.region = region
         self.endpoint = endpoint
         self._client = S3Client(region, endpoint)

s3torchconnector/src/s3torchconnector/s3iterable_dataset.py

@@ -1,7 +1,7 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
 from functools import partial
-from typing import Iterator, Any, Union, Iterable, Callable
+from typing import Iterator, Any, Union, Iterable, Callable, Optional
 import logging
 
 import torch.utils.data
@@ -29,7 +29,7 @@ def __init__(
         self,
         region: str,
         get_dataset_objects: Callable[[S3Client], Iterable[S3BucketKeyData]],
-        endpoint: str = None,
+        endpoint: Optional[str] = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         self._get_dataset_objects = get_dataset_objects
@@ -52,7 +52,7 @@ def from_objects(
         object_uris: Union[str, Iterable[str]],
         *,
         region: str,
-        endpoint: str = None,
+        endpoint: Optional[str] = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         """Returns an instance of S3IterableDataset using the S3 URI(s) provided.
@@ -83,7 +83,7 @@ def from_prefix(
         s3_uri: str,
         *,
         region: str,
-        endpoint: str = None,
+        endpoint: Optional[str] = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         """Returns an instance of S3IterableDataset using the S3 URI provided.

s3torchconnector/src/s3torchconnector/s3map_dataset.py

@@ -1,7 +1,7 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  // SPDX-License-Identifier: BSD
 from functools import partial
-from typing import List, Any, Callable, Iterable, Union
+from typing import List, Any, Callable, Iterable, Union, Optional
 import logging
 
 import torch.utils.data
@@ -30,15 +30,15 @@ def __init__(
         self,
         region: str,
         get_dataset_objects: Callable[[S3Client], Iterable[S3BucketKeyData]],
-        endpoint: str = None,
+        endpoint: Optional[str] = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         self._get_dataset_objects = get_dataset_objects
         self._transform = transform
         self._region = region
         self._endpoint = endpoint
         self._client = None
-        self._bucket_key_pairs = None
+        self._bucket_key_pairs: Optional[List[S3BucketKeyData]] = None
 
     @property
     def region(self):
@@ -52,6 +52,7 @@ def endpoint(self):
     def _dataset_bucket_key_pairs(self) -> List[S3BucketKeyData]:
         if self._bucket_key_pairs is None:
             self._bucket_key_pairs = list(self._get_dataset_objects(self._get_client()))
+        assert self._bucket_key_pairs is not None
         return self._bucket_key_pairs
 
     @classmethod
@@ -60,7 +61,7 @@ def from_objects(
         object_uris: Union[str, Iterable[str]],
         *,
         region: str,
-        endpoint: str = None,
+        endpoint: Optional[str] = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         """Returns an instance of S3MapDataset using the S3 URI(s) provided.
@@ -91,7 +92,7 @@ def from_prefix(
         s3_uri: str,
         *,
         region: str,
-        endpoint: str = None,
+        endpoint: Optional[str] = None,
         transform: Callable[[S3Reader], Any] = identity,
     ):
         """Returns an instance of S3MapDataset using the S3 URI provided.

s3torchconnector/src/s3torchconnector/s3reader.py

@@ -4,7 +4,7 @@
 import io
 from functools import cached_property
 from io import SEEK_CUR, SEEK_END, SEEK_SET
-from typing import Callable, Optional
+from typing import Callable, Optional, Iterable, Iterator
 
 from s3torchconnectorclient._mountpoint_s3_client import ObjectInfo, GetObjectStream
 
@@ -16,18 +16,18 @@ def __init__(
         self,
         bucket: str,
         key: str,
-        get_object_info: Callable[[], ObjectInfo] = None,
-        get_stream: Callable[[], GetObjectStream] = None,
+        get_object_info: Callable[[], ObjectInfo],
+        get_stream: Callable[[], GetObjectStream],
     ):
         if not bucket:
             raise ValueError("Bucket should be specified")
         self._bucket = bucket
         self._key = key
         self._get_object_info = get_object_info
         self._get_stream = get_stream
-        self._stream = None
+        self._stream: Optional[Iterator[bytes]] = None
         self._buffer = io.BytesIO()
-        self._size = None
+        self._size: Optional[int] = None
         # Invariant: _position == _buffer._tell() unless _position_at_end()
         self._position = 0
 
@@ -77,6 +77,7 @@ def read(self, size: Optional[int] = None) -> bytes:
             return b""
 
         self.prefetch()
+        assert self._stream is not None
         cur_pos = self._position
         if size is None or size < 0:
             # Special case read() all to use O(n) algorithm
@@ -138,6 +139,7 @@ def seek(self, offset: int, whence: int = SEEK_SET, /) -> int:
 
     def _prefetch_to_offset(self, offset: int) -> None:
         self.prefetch()
+        assert self._stream is not None
         buf_size = self._buffer.seek(0, SEEK_END)
         try:
             while offset > buf_size: