diff --git a/lib/dt_shell/environments.py b/lib/dt_shell/environments.py index a030037..b6e3c62 100644 --- a/lib/dt_shell/environments.py +++ b/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_tool, 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_tool(interpreter_fpath) # install dependencies cache: InstalledDependenciesDatabase = InstalledDependenciesDatabase.load(shell.profile) diff --git a/lib/dt_shell/utils.py b/lib/dt_shell/utils.py index 896cd55..43a561a 100644 --- a/lib/dt_shell/utils.py +++ b/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_tool(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_tool(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"])