s3_management: Fix pyright / mypy errors add mypy linter (#6904)

I was observing these errors in my editor and they were annoying me so I
fixed them. I don't expect any functional changes to come out of this.

---------

Signed-off-by: Eli Uriegas <[email protected]>

Author: seemethere

.lintrunner.toml

@@ -46,6 +46,7 @@ include_patterns = [
     'tools/**/*.pyi',
     'torchci/**/*.py',
     'torchci/**/*.pyi',
+    's3_management/*.py',
 ]
 exclude_patterns = [
     'aws/lambda/servicelab-ingestor/**',

s3_management/backup_conda.py

@@ -7,11 +7,11 @@
 import argparse
 import hashlib
 import os
-import urllib
+import urllib.request
 from typing import List, Optional
 
-import boto3
-import conda.api
+import boto3  # type: ignore[import-untyped]
+import conda.api  # type: ignore[import-not-found]
 
 
 S3 = boto3.resource("s3")

s3_management/manage.py

@@ -11,10 +11,10 @@
 from datetime import datetime
 from os import makedirs, path
 from re import match, search, sub
-from typing import Dict, Iterable, List, Optional, Set, Type, TypeVar
+from typing import Dict, Iterable, List, Optional, Set, Type, TypeVar, Union
 
-import boto3
-import botocore
+import boto3  # type: ignore[import]
+import botocore  # type: ignore[import]
 from packaging.version import InvalidVersion, parse as _parse_version, Version
 
 
@@ -240,13 +240,13 @@ def __lt__(self, other):
 
 def safe_parse_version(ver_str: str) -> Version:
     try:
-        return _parse_version(ver_str)
+        return _parse_version(ver_str)  # type: ignore[return-value]
     except InvalidVersion:
         return Version("0.0.0")
 
 
 class S3Index:
-    def __init__(self: S3IndexType, objects: List[S3Object], prefix: str) -> None:
+    def __init__(self, objects: List[S3Object], prefix: str) -> None:
         self.objects = objects
         self.prefix = prefix.rstrip("/")
         self.html_name = "index.html"
@@ -256,7 +256,7 @@ def __init__(self: S3IndexType, objects: List[S3Object], prefix: str) -> None:
             path.dirname(obj.key) for obj in objects if path.dirname != prefix
         }
 
-    def nightly_packages_to_show(self: S3IndexType) -> List[S3Object]:
+    def nightly_packages_to_show(self) -> List[S3Object]:
         """Finding packages to show based on a threshold we specify
 
         Basically takes our S3 packages, normalizes the version for easier
@@ -326,7 +326,7 @@ def get_package_names(self, subdir: Optional[str] = None) -> List[str]:
             {self.obj_to_package_name(obj) for obj in self.gen_file_list(subdir)}
         )
 
-    def normalize_package_version(self: S3IndexType, obj: S3Object) -> str:
+    def normalize_package_version(self, obj: S3Object) -> str:
         # removes the GPU specifier from the package name as well as
         # unnecessary things like the file extension, architecture name, etc.
         return sub(r"%2B.*", "", "-".join(path.basename(obj.key).split("-")[:2]))
@@ -498,7 +498,7 @@ def compute_sha256(self) -> None:
             )
 
     @classmethod
-    def has_public_read(cls: Type[S3IndexType], key: str) -> bool:
+    def has_public_read(cls, key: str) -> bool:
         def is_all_users_group(o) -> bool:
             return (
                 o.get("Grantee", {}).get("URI")
@@ -512,11 +512,11 @@ def can_read(o) -> bool:
         return any(is_all_users_group(x) and can_read(x) for x in acl_grants)
 
     @classmethod
-    def grant_public_read(cls: Type[S3IndexType], key: str) -> None:
+    def grant_public_read(cls, key: str) -> None:
         CLIENT.put_object_acl(Bucket=BUCKET.name, Key=key, ACL="public-read")
 
     @classmethod
-    def fetch_object_names(cls: Type[S3IndexType], prefix: str) -> List[str]:
+    def fetch_object_names(cls, prefix: str) -> List[str]:
         obj_names = []
         for obj in BUCKET.objects.filter(Prefix=prefix):
             is_acceptable = any(
@@ -531,7 +531,7 @@ def fetch_object_names(cls: Type[S3IndexType], prefix: str) -> List[str]:
             obj_names.append(obj.key)
         return obj_names
 
-    def fetch_metadata(self: S3IndexType) -> None:
+    def fetch_metadata(self) -> None:
         # Add PEP 503-compatible hashes to URLs to allow clients to avoid spurious downloads, if possible.
         regex_multipart_upload = r"^[A-Za-z0-9+/=]+=-[0-9]+$"
         with concurrent.futures.ThreadPoolExecutor(max_workers=6) as executor:
@@ -565,17 +565,17 @@ def fetch_metadata(self: S3IndexType) -> None:
                 if size := response.get("ContentLength"):
                     self.objects[idx].size = int(size)
 
-    def fetch_pep658(self: S3IndexType) -> None:
+    def fetch_pep658(self) -> None:
         def _fetch_metadata(key: str) -> str:
             try:
                 response = CLIENT.head_object(
                     Bucket=BUCKET.name, Key=f"{key}.metadata", ChecksumMode="Enabled"
                 )
                 sha256 = base64.b64decode(response.get("ChecksumSHA256")).hex()
                 return sha256
-            except botocore.exceptions.ClientError as e:
+            except botocore.exceptions.ClientError as e:  # type: ignore[attr-defined]
                 if e.response["Error"]["Code"] == "404":
-                    return None
+                    return ""
                 raise
 
         with concurrent.futures.ThreadPoolExecutor(max_workers=6) as executor:
@@ -592,9 +592,7 @@ def _fetch_metadata(key: str) -> str:
                     self.objects[idx].pep658 = response
 
     @classmethod
-    def from_S3(
-        cls: Type[S3IndexType], prefix: str, with_metadata: bool = True
-    ) -> S3IndexType:
+    def from_S3(cls, prefix: str, with_metadata: bool = True) -> "S3Index":
         prefix = prefix.rstrip("/")
         obj_names = cls.fetch_object_names(prefix)
 
@@ -622,7 +620,7 @@ def sanitize_key(key: str) -> str:
         return rc
 
     @classmethod
-    def undelete_prefix(cls: Type[S3IndexType], prefix: str) -> None:
+    def undelete_prefix(cls, prefix: str) -> None:
         paginator = CLIENT.get_paginator("list_object_versions")
         for page in paginator.paginate(Bucket=BUCKET.name, Prefix=prefix):
             for obj in page.get("DeleteMarkers", []):

s3_management/update_dependencies.py

@@ -1,7 +1,7 @@
 import re
 from typing import Dict, List
 
-import boto3
+import boto3  # type: ignore[import-untyped]
 
 
 S3 = boto3.resource("s3")
@@ -225,7 +225,7 @@ def upload_missing_whls(
 
     pypi_latest_packages = get_wheels_of_version(pypi_idx, selected_version)
 
-    download_latest_packages = []
+    download_latest_packages: Dict[str, str] = {}
     if not only_pypi:
         download_idx = parse_simple_idx(
             f"https://download.pytorch.org/{prefix}/{pkg_name}"