fix: streamline pip installation process in virtual environments

Author: liammcalpineduckietown

lib/dt_shell/environments.py

@@ -14,7 +14,7 @@
 from .constants import SHELL_LIB_DIR, SHELL_REQUIREMENTS_LIST, DTShellConstants
 from .database.utils import InstalledDependenciesDatabase
 from .logging import dts_print
-from .utils import pip_install, replace_spaces, print_debug_info, pretty_json
+from .utils import install_pip, pip_install, replace_spaces, print_debug_info, pretty_json
 
 
 class ShellCommandEnvironmentAbs(metaclass=ABCMeta):
@@ -110,16 +110,7 @@ def execute(self, shell, _: List[str]):
                 with_pip=False,
                 prompt="dts"
             )
-            # install pip
-            get_pip_fpath: str = os.path.join(SHELL_LIB_DIR, "assets", "get-pip.py")
-            assert os.path.exists(get_pip_fpath)
-            logger.info(f"Installing pip...")
-            try:
-                subprocess.check_output([interpreter_fpath, get_pip_fpath], stderr=subprocess.PIPE)
-            except subprocess.CalledProcessError as e:
-                # TODO: test this failure case on purpose
-                msg: str = "An error occurred while installing pip in the virtual environment"
-                raise ShellInitException(msg, stdout=e.stdout, stderr=e.stderr)
+            install_pip(interpreter_fpath)
 
         # install dependencies
         cache: InstalledDependenciesDatabase = InstalledDependenciesDatabase.load(shell.profile)

lib/dt_shell/utils.py

@@ -27,6 +27,7 @@
 from .exceptions import ShellInitException, RunCommandException
 
 NOTSET = object()
+MAX_PIP_INSTALL_ATTEMPTS = 2
 
 
 cli_style = Style([
@@ -235,19 +236,38 @@ class DebugInfo:
     name2versions: Dict[str, Union[str, Dict[str, str]]] = {}
 
 
-def pip_install(interpreter: str, requirements: str):
-    run = subprocess.check_call if logger.level <= logging.DEBUG else subprocess.check_output
+def install_pip(interpreter: str):
+    get_pip_fpath: str = os.path.join(SHELL_LIB_DIR, "assets", "get-pip.py")
+    if not os.path.exists(get_pip_fpath):
+        msg = f"Required file for pip installation not found: {get_pip_fpath}"
+        raise ShellInitException(msg)
+    logger.info("Installing pip...")
     try:
-        run(
-            [interpreter, "-m", "pip", "install", "-r", requirements],
-            stderr=subprocess.STDOUT,
-            env={}
-        )
+        subprocess.check_output([interpreter, get_pip_fpath], stderr=subprocess.PIPE)
     except subprocess.CalledProcessError as e:
-        msg: str = "An error occurred while installing python dependencies"
+        # TODO: test this failure case on purpose
+        msg: str = "An error occurred while installing pip in the virtual environment"
         raise ShellInitException(msg, stdout=e.stdout, stderr=e.stderr)
 
 
+def pip_install(interpreter: str, requirements: str):
+    run = subprocess.check_call if logger.level <= logging.DEBUG else subprocess.check_output
+    for attempt in range(MAX_PIP_INSTALL_ATTEMPTS):
+        try:
+            run(
+                [interpreter, "-m", "pip", "install", "-r", requirements],
+                stderr=subprocess.STDOUT,
+                env={}
+            )
+            break
+        except subprocess.CalledProcessError as e:
+            error = e.stdout.decode("utf-8", errors="replace") if e.stdout else ""
+            if attempt == MAX_PIP_INSTALL_ATTEMPTS - 1 or "No module named pip" not in error:
+                msg: str = "An error occurred while installing python dependencies"
+                raise ShellInitException(msg, stdout=e.stdout, stderr=e.stderr)
+            install_pip(interpreter)
+
+
 def indent_block(s: str, indent_len: int = 4) -> str:
     space: str = " " * indent_len
     return space + f"\n{space}".join(s.splitlines() if s is not None else ["None"])