Add type: ignore annotations in various places to make mypy pass (#168)

* Add `type: ignore` annotations in various places to make mypy pass

* Install lightning extension for type checks

* Add note about how to run mypy without lightning installed

---------

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

Author: muddyfish

.github/workflows/python-checks.yml

@@ -97,7 +97,7 @@ jobs:
           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
+          python -m pip install ./s3torchconnector[lightning]
 
       - name: Lint with flake8
         run: |

doc/DEVELOPMENT.md

@@ -75,6 +75,12 @@ mypy s3torchconnectorclient/python/src
 ```
  to lint.
 
+To run mypy without `lightning` installed, run
+```bash
+mypy s3torchconnector/src --exclude s3torchconnector/src/s3torchconnector/lightning
+mypy s3torchconnectorclient/python/src
+```
+
 ### Debugging
 
 Either a Python or GDB style debugger will be useful here.

s3torchconnector/src/s3torchconnector/lightning/s3_lightning_checkpoint.py

@@ -24,7 +24,8 @@ def __init__(self, region: str):
     def save_checkpoint(
         self,
         checkpoint: Dict[str, Any],
-        path: str,
+        # We only support `str` arguments for `path`, as `Path` is explicitly for local filesystems
+        path: str,  # type: ignore
         storage_options: Optional[Any] = None,
     ) -> None:
         """Save model/training states as a checkpoint file through state-dump and upload to S3.
@@ -41,7 +42,8 @@ def save_checkpoint(
 
     def load_checkpoint(
         self,
-        path: str,
+        # We only support `str` arguments for `path`, as `Path` is explicitly for local filesystems
+        path: str,  # type: ignore
         map_location: Optional[Any] = None,
     ) -> Dict[str, Any]:
         """Load checkpoint from an S3 location when resuming or loading ckpt for test/validate/predict stages.
@@ -59,9 +61,15 @@ def load_checkpoint(
         self._validate_path(path)
         bucket, key = parse_s3_uri(path)
         s3reader = self._client.get_object(bucket, key)
-        return torch.load(s3reader, map_location)
+        # FIXME - io.BufferedIOBase and typing.IO aren't compatible
+        #  See https://github.com/python/typeshed/issues/6077
+        return torch.load(s3reader, map_location)  # type: ignore
 
-    def remove_checkpoint(self, path: str) -> None:
+    def remove_checkpoint(
+        self,
+        # We only support `str` arguments for `path`, as `Path` is explicitly for local filesystems
+        path: str,  # type: ignore
+    ) -> None:
         """Remove checkpoint file from the S3 uri.
 
         Args: