Type-annotate wandb/sdk/internal/update.py (#3432)

* type-annotate sdk.internal.update.py

mypy.ini

@@ -223,7 +223,7 @@ disallow_untyped_defs = False
 disallow_untyped_calls = False
 
 [mypy-wandb.sdk.internal.update]
-ignore_errors = True
+ignore_errors = False
 
 [mypy-wandb.sdk.internal.profiler]
 disallow_untyped_calls = False

wandb/sdk/internal/update.py

@@ -1,17 +1,22 @@
+from typing import Dict, Optional, Tuple
+
 from pkg_resources import parse_version
 import requests
 import wandb
 
 
-def _find_available(current_version):
-    pypi_url = "https://pypi.org/pypi/%s/json" % wandb._wandb_module
+def _find_available(
+    current_version: str,
+) -> Optional[Tuple[str, bool, bool, bool, Optional[str]]]:
+    pypi_url = f"https://pypi.org/pypi/{wandb._wandb_module}/json"
 
     yanked_dict = {}
     try:
+        # raise Exception("test")
         async_requests_get = wandb.util.async_call(requests.get, timeout=5)
         data, thread = async_requests_get(pypi_url, timeout=3)
         if not data or isinstance(data, Exception):
-            return
+            return None
         data = data.json()
         latest_version = data["info"]["version"]
         release_list = data["releases"].keys()
@@ -23,7 +28,7 @@ def _find_available(current_version):
                     yanked_dict[version] = yanked_reason
     except Exception:
         # Any issues whatsoever, just skip the latest version check.
-        return
+        return None
 
     # Return if no update is available
     pip_prerelease = False
@@ -44,56 +49,57 @@ def _find_available(current_version):
     if parse_version(latest_version) <= parsed_current_version:
         # pre-releases are not included in latest_version
         # so if we are currently running a pre-release we check more
-        if not parsed_current_version.is_prerelease:
-            return
+        if not parsed_current_version.is_prerelease:  # type: ignore
+            return None
         # Candidates are pre-releases with the same base_version
         release_list = map(parse_version, release_list)
         release_list = filter(lambda v: v.is_prerelease, release_list)
         release_list = filter(
-            lambda v: v.base_version == parsed_current_version.base_version,
+            lambda v: v.base_version == parsed_current_version.base_version,  # type: ignore
             release_list,
         )
         release_list = sorted(release_list)
         if not release_list:
-            return
+            return None
 
         parsed_latest_version = release_list[-1]
         if parsed_latest_version <= parsed_current_version:
-            return
+            return None
         latest_version = str(parsed_latest_version)
         pip_prerelease = True
 
     return latest_version, pip_prerelease, deleted, yanked, yanked_reason
 
 
-def check_available(current_version):
+def check_available(current_version: str) -> Optional[Dict[str, Optional[str]]]:
     package_info = _find_available(current_version)
     if not package_info:
-        return
+        return None
+
+    wandb_module_name = wandb._wandb_module
 
     latest_version, pip_prerelease, deleted, yanked, yanked_reason = package_info
     upgrade_message = (
         "%s version %s is available!  To upgrade, please run:\n"
         " $ pip install %s --upgrade%s"
         % (
-            wandb._wandb_module,
+            wandb_module_name,
             latest_version,
-            wandb._wandb_module,
+            wandb_module_name,
             " --pre" if pip_prerelease else "",
         )
     )
-
     delete_message = None
     if deleted:
         delete_message = "{} version {} has been retired!  Please upgrade.".format(
-            wandb._wandb_module,
+            wandb_module_name,
             current_version,
         )
     yank_message = None
     if yanked:
         reason_message = "(%s)  " % yanked_reason if yanked_reason else ""
         yank_message = "{} version {} has been recalled!  {}Please upgrade.".format(
-            wandb._wandb_module,
+            wandb_module_name,
             current_version,
             reason_message,
         )