From e85bd8beba95238f28fe94d3a69f1fa6d6743f73 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Tue, 5 Dec 2023 07:45:52 +0900
Subject: [PATCH] fix: use py3.8 compatible annotations and cleanup imports

---
 .../tools/wheel_installer/wheel.py            | 24 +++++++++----------
 .../tools/wheel_installer/wheel_installer.py  |  7 ++----
 2 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/python/pip_install/tools/wheel_installer/wheel.py b/python/pip_install/tools/wheel_installer/wheel.py
index 73f83eb8ff..438a7bec25 100644
--- a/python/pip_install/tools/wheel_installer/wheel.py
+++ b/python/pip_install/tools/wheel_installer/wheel.py
@@ -14,12 +14,12 @@
 
 """Utility class to inspect an extracted wheel directory"""
 
-from __future__ import annotations
-
 import email
 from collections import defaultdict
 from dataclasses import dataclass, field
 from enum import Enum, unique
+from pathlib import Path
+from typing import Dict, List, Optional, Set, Tuple
 
 import installer
 import pkg_resources
@@ -33,7 +33,7 @@ class OS(Enum):
     windows = 3
 
     @staticmethod
-    def from_tag(tag: str) -> OS:
+    def from_tag(tag: str) -> "OS":
         if tag.startswith("linux"):
             return OS.linux
         elif tag.startswith("manylinux"):
@@ -89,7 +89,7 @@ def __str__(self) -> str:
         return self.os.name.lower() + "_" + self.arch.name.lower()
 
     @classmethod
-    def from_tag(cls, tag: str) -> Platform:
+    def from_tag(cls, tag: str) -> "Platform":
         return cls(
             os=OS.from_tag(tag),
             arch=Arch.from_tag(tag),
@@ -161,10 +161,10 @@ def _default_select():
 
 @dataclass
 class Deps:
-    deps: set[str] = field(default_factory=set)
-    select: dict[Platform, set[str]] = field(default_factory=_default_select)
+    deps: Set[str] = field(default_factory=set)
+    select: Dict[Platform, Set[str]] = field(default_factory=_default_select)
 
-    def add(self, dep: str, platform: Platform | None = None):
+    def add(self, dep: str, platform: Optional[Platform] = None):
         if platform:
             self.select[platform].add(dep)
         else:
@@ -172,7 +172,7 @@ def add(self, dep: str, platform: Platform | None = None):
             for p, deps in self.select.items():
                 self.select[p] = deps - {dep}
 
-    def build(self, all_platforms: list[Platform]) -> Deps:
+    def build(self, all_platforms: List[Platform]) -> "Deps":
         # Move deps to common deps if they are present for all platforms
         common_deps = None
         for plat in all_platforms:
@@ -192,7 +192,7 @@ def build(self, all_platforms: list[Platform]) -> Deps:
 class Wheel:
     """Representation of the compressed .whl file"""
 
-    def __init__(self, path: str):
+    def __init__(self, path: Path):
         self._path = path
 
     @property
@@ -217,13 +217,13 @@ def version(self) -> str:
         # TODO Also available as installer.sources.WheelSource.version
         return str(self.metadata["Version"])
 
-    def entry_points(self) -> dict[str, tuple[str, str]]:
+    def entry_points(self) -> Dict[str, Tuple[str, str]]:
         """Returns the entrypoints defined in the current wheel
 
         See https://packaging.python.org/specifications/entry-points/ for more info
 
         Returns:
-            dict[str, tuple[str, str]]: A mapping of the entry point's name to it's module and attribute
+            Dict[str, Tuple[str, str]]: A mapping of the entry point's name to it's module and attribute
         """
         with installer.sources.WheelFile.open(self.path) as wheel_source:
             if "entry_points.txt" not in wheel_source.dist_info_filenames:
@@ -239,7 +239,7 @@ def entry_points(self) -> dict[str, tuple[str, str]]:
             return entry_points_mapping
 
     def dependencies(
-        self, extras_requested: set[str] = None, platforms=list[Platform]
+        self, extras_requested: Set[str] = None, platforms=List[Platform]
     ) -> Deps:
         # NOTE @aignas 2023-12-04: if the wheel is a platform specific wheel, we only include deps for that platform
         _, _, platform_tag = self._path.name.rpartition("-")
diff --git a/python/pip_install/tools/wheel_installer/wheel_installer.py b/python/pip_install/tools/wheel_installer/wheel_installer.py
index 2460f5c451..ad79c5953a 100644
--- a/python/pip_install/tools/wheel_installer/wheel_installer.py
+++ b/python/pip_install/tools/wheel_installer/wheel_installer.py
@@ -14,19 +14,16 @@
 
 """Build and/or fetch a single wheel based on the requirement passed in"""
 
-import argparse
 import errno
 import glob
 import json
 import os
 import re
-import shutil
 import subprocess
 import sys
-import textwrap
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import Dict, Iterable, List, Optional, Set, Tuple
+from typing import Dict, List, Optional, Set, Tuple
 
 from pip._vendor.packaging.utils import canonicalize_name
 
@@ -108,7 +105,7 @@ def _extract_wheel(
     wheel_file: str,
     extras: Dict[str, Set[str]],
     enable_implicit_namespace_pkgs: bool,
-    platforms: list[wheel.Platform],
+    platforms: List[wheel.Platform],
     installation_dir: Path = Path("."),
 ) -> None:
     """Extracts wheel into given directory and creates py_library and filegroup targets.