installer: use venv and improve error handling

.github/workflows/installer.yml

@@ -49,6 +49,13 @@ jobs:
         shell: bash
         run: python install-poetry.py -y
 
+      - name: Upload Failure Log
+        uses: actions/upload-artifact@v2
+        if: failure()
+        with:
+          name: poetry-installer-error.log
+          path: poetry-installer-error-*.log
+
       - name: Verify Installation
         shell: bash
         run: |

install-poetry.py

@@ -3,14 +3,15 @@
 
 It does, in order:
 
-  - Downloads the virtualenv package to a temporary directory and add it to sys.path.
-  - Creates a virtual environment in the correct OS data dir which will be
+  - Creates a virtual environment using venv (or virtualenv zipapp) in the correct OS data dir which will be
       - `%APPDATA%\\pypoetry` on Windows
       -  ~/Library/Application Support/pypoetry on MacOS
       - `${XDG_DATA_HOME}/pypoetry` (or `~/.local/share/pypoetry` if it's not set) on UNIX systems
       - In `${POETRY_HOME}` if it's set.
   - Installs the latest or given version of Poetry inside this virtual environment.
   - Installs a `poetry` script in the Python user directory (or `${POETRY_HOME/bin}` if `POETRY_HOME` is set).
+  - On failure, the error log is written to poetry-installer-error-*.log and any previously existing environment
+    is restored.
 """
 
 import argparse
@@ -219,21 +220,6 @@ def _get_win_folder_with_ctypes(csidl_name):
         _get_win_folder = _get_win_folder_from_registry
 
 
-@contextmanager
-def temporary_directory(*args, **kwargs):
-    try:
-        from tempfile import TemporaryDirectory
-    except ImportError:
-        name = tempfile.mkdtemp(*args, **kwargs)
-
-        yield name
-
-        shutil.rmtree(name)
-    else:
-        with TemporaryDirectory(*args, **kwargs) as name:
-            yield name
-
-
 PRE_MESSAGE = """# Welcome to {poetry}!
 
 This will download and install the latest version of {poetry},
@@ -277,6 +263,83 @@ def temporary_directory(*args, **kwargs):
 POST_MESSAGE_CONFIGURE_WINDOWS = """"""
 
 
+class PoetryInstallationError(RuntimeError):
+    def __init__(self, return_code: int = 0, log: Optional[str] = None):
+        super(PoetryInstallationError, self).__init__()
+        self.return_code = return_code
+        self.log = log
+
+
+class VirtualEnvironment:
+    def __init__(self, path: Path) -> None:
+        self._path = path
+        # str is required for compatibility with subprocess run on CPython <= 3.7 on Windows
+        self._python = str(
+            self._path.joinpath("Scripts/python.exe" if WINDOWS else "bin/python")
+        )
+
+    @property
+    def path(self):
+        return self._path
+
+    @classmethod
+    def make(cls, target: Path) -> "VirtualEnvironment":
+        try:
+            import venv
+
+            builder = venv.EnvBuilder(clear=True, with_pip=True, symlinks=False)
+            builder.ensure_directories(target)
+            builder.create(target)
+        except ImportError:
+            # fallback to using virtualenv package if venv is not available, eg: ubuntu
+            python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+            virtualenv_bootstrap_url = (
+                f"https://bootstrap.pypa.io/virtualenv/{python_version}/virtualenv.pyz"
+            )
+
+            with tempfile.TemporaryDirectory(prefix="poetry-installer") as temp_dir:
+                virtualenv_pyz = Path(temp_dir) / "virtualenv.pyz"
+                request = Request(
+                    virtualenv_bootstrap_url, headers={"User-Agent": "Python Poetry"}
+                )
+                virtualenv_pyz.write_bytes(urlopen(request).read())
+                cls.run(
+                    sys.executable, virtualenv_pyz, "--clear", "--always-copy", target
+                )
+
+        # We add a special file so that Poetry can detect
+        # its own virtual environment
+        target.joinpath("poetry_env").touch()
+
+        env = cls(target)
+
+        # we do this here to ensure that outdated system default pip does not trigger older bugs
+        env.pip("install", "--disable-pip-version-check", "--upgrade", "pip")
+
+        return env
+
+    @staticmethod
+    def run(*args, **kwargs) -> subprocess.CompletedProcess:
+        completed_process = subprocess.run(
+            args,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            **kwargs,
+        )
+        if completed_process.returncode != 0:
+            raise PoetryInstallationError(
+                return_code=completed_process.returncode,
+                log=completed_process.stdout.decode(),
+            )
+        return completed_process
+
+    def python(self, *args, **kwargs) -> subprocess.CompletedProcess:
+        return self.run(self._python, *args, **kwargs)
+
+    def pip(self, *args, **kwargs) -> subprocess.CompletedProcess:
+        return self.python("-m", "pip", "--isolated", *args, **kwargs)
+
+
 class Cursor:
     def __init__(self) -> None:
         self._output = sys.stdout
@@ -439,12 +502,10 @@ def _is_self_upgrade_supported(x):
         try:
             self.install(version)
         except subprocess.CalledProcessError as e:
-            print(
-                colorize("error", f"\nAn error has occurred: {e}\n{e.stdout.decode()}")
+            raise PoetryInstallationError(
+                return_code=e.returncode, log=e.output.decode()
             )
 
-            return e.returncode
-
         self._write("")
         self.display_post_message(version)
 
@@ -460,21 +521,13 @@ def install(self, version, upgrade=False):
             )
         )
 
-        env_path = self.make_env(version)
-        self.install_poetry(version, env_path)
-        self.make_bin(version)
+        with self.make_env(version) as env:
+            self.install_poetry(version, env)
+            self.make_bin(version, env)
+            self._data_dir.joinpath("VERSION").write_text(version)
+            self._install_comment(version, "Done")
 
-        self._overwrite(
-            "Installing {} ({}): {}".format(
-                colorize("info", "Poetry"),
-                colorize("b", version),
-                colorize("success", "Done"),
-            )
-        )
-
-        self._data_dir.joinpath("VERSION").write_text(version)
-
-        return 0
+            return 0
 
     def uninstall(self) -> int:
         if not self._data_dir.exists():
@@ -504,81 +557,70 @@ def uninstall(self) -> int:
 
         return 0
 
-    def make_env(self, version: str) -> Path:
+    def _install_comment(self, version: str, message: str):
         self._overwrite(
             "Installing {} ({}): {}".format(
                 colorize("info", "Poetry"),
                 colorize("b", version),
-                colorize("comment", "Creating environment"),
+                colorize("comment", message),
             )
         )
 
+    @contextmanager
+    def make_env(self, version: str) -> VirtualEnvironment:
         env_path = self._data_dir.joinpath("venv")
+        env_path_saved = env_path.with_suffix(".save")
 
-        with temporary_directory() as tmp_dir:
-            subprocess.run(
-                [sys.executable, "-m", "pip", "install", "virtualenv", "-t", tmp_dir],
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                check=True,
-            )
-
-            sys.path.insert(0, tmp_dir)
-
-            import virtualenv
+        if env_path.exists():
+            self._install_comment(version, "Saving existing environment")
+            if env_path_saved.exists():
+                shutil.rmtree(env_path_saved)
+            shutil.move(env_path, env_path_saved)
 
-            virtualenv.cli_run([str(env_path), "--clear"])
-
-        # We add a special file so that Poetry can detect
-        # its own virtual environment
-        env_path.joinpath("poetry_env").touch()
+        try:
+            self._install_comment(version, "Creating environment")
+            yield VirtualEnvironment.make(env_path)
+        except Exception as e:  # noqa
+            if env_path.exists():
+                self._install_comment(
+                    version, "An error occurred. Removing partial environment."
+                )
+                shutil.rmtree(env_path)
 
-        return env_path
+            if env_path_saved.exists():
+                self._install_comment(
+                    version, "Restoring previously saved environment."
+                )
+                shutil.move(env_path_saved, env_path)
 
-    def make_bin(self, version: str) -> None:
-        self._overwrite(
-            "Installing {} ({}): {}".format(
-                colorize("info", "Poetry"),
-                colorize("b", version),
-                colorize("comment", "Creating script"),
-            )
-        )
+            raise e
+        else:
+            if env_path_saved.exists():
+                shutil.rmtree(env_path_saved, ignore_errors=True)
 
+    def make_bin(self, version: str, env: VirtualEnvironment) -> None:
+        self._install_comment(version, "Creating script")
         self._bin_dir.mkdir(parents=True, exist_ok=True)
 
         script = "poetry"
-        target_script = "venv/bin/poetry"
+        script_bin = "bin"
         if WINDOWS:
             script = "poetry.exe"
-            target_script = "venv/Scripts/poetry.exe"
+            script_bin = "Scripts"
+        target_script = env.path.joinpath(script_bin, script)
 
         if self._bin_dir.joinpath(script).exists():
             self._bin_dir.joinpath(script).unlink()
 
         try:
-            self._bin_dir.joinpath(script).symlink_to(
-                self._data_dir.joinpath(target_script)
-            )
+            self._bin_dir.joinpath(script).symlink_to(target_script)
         except OSError:
             # This can happen if the user
             # does not have the correct permission on Windows
-            shutil.copy(
-                self._data_dir.joinpath(target_script), self._bin_dir.joinpath(script)
-            )
+            shutil.copy(target_script, self._bin_dir.joinpath(script))
 
-    def install_poetry(self, version: str, env_path: Path) -> None:
-        self._overwrite(
-            "Installing {} ({}): {}".format(
-                colorize("info", "Poetry"),
-                colorize("b", version),
-                colorize("comment", "Installing Poetry"),
-            )
-        )
-
-        if WINDOWS:
-            python = env_path.joinpath("Scripts/python.exe")
-        else:
-            python = env_path.joinpath("bin/python")
+    def install_poetry(self, version: str, env: VirtualEnvironment) -> None:
+        self._install_comment(version, "Installing Poetry")
 
         if self._git:
             specification = "git+" + version
@@ -587,12 +629,7 @@ def install_poetry(self, version: str, env_path: Path) -> None:
         else:
             specification = f"poetry=={version}"
 
-        subprocess.run(
-            [str(python), "-m", "pip", "install", specification],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            check=True,
-        )
+        env.pip("install", specification)
 
     def display_pre_message(self) -> None:
         kwargs = {
@@ -831,7 +868,25 @@ def main():
     if args.uninstall or string_to_bool(os.getenv("POETRY_UNINSTALL", "0")):
         return installer.uninstall()
 
-    return installer.run()
+    try:
+        return installer.run()
+    except PoetryInstallationError as e:
+        installer._write(colorize("error", "Poetry installation failed."))  # noqa
+
+        if e.log is not None:
+            import traceback
+
+            _, path = tempfile.mkstemp(
+                suffix=".log",
+                prefix="poetry-installer-error-",
+                dir=str(Path.cwd()),
+                text=True,
+            )
+            installer._write(colorize("error", f"See {path} for error logs."))  # noqa
+            text = f"{e.log}\nTraceback:\n\n{''.join(traceback.format_tb(e.__traceback__))}"
+            Path(path).write_text(text)
+
+        return e.return_code
 
 
 if __name__ == "__main__":