Improve the self update command and environment detection

poetry/console/commands/self/update.py

@@ -4,6 +4,7 @@
 import os
 import re
 import shutil
+import site
 import stat
 import subprocess
 import sys
@@ -16,6 +17,8 @@
 from cleo import option
 
 from poetry.core.packages import Dependency
+from poetry.utils._compat import PY2
+from poetry.utils._compat import Path
 
 from ..command import Command
 
@@ -60,6 +63,10 @@ class SelfUpdateCommand(Command):
     REPOSITORY_URL = "https://github.com/python-poetry/poetry"
     BASE_URL = REPOSITORY_URL + "/releases/download"
 
+    _data_dir = None
+    _bin_dir = None
+    _pool = None
+
     @property
     def home(self):
         from poetry.utils._compat import Path
@@ -78,18 +85,75 @@ def lib(self):
     def lib_backup(self):
         return self.home / "lib-backup"
 
+    @property
+    def data_dir(self):  # type: () -> Path
+        if self._data_dir is not None:
+            return self._data_dir
+
+        from poetry.locations import data_dir
+
+        self._data_dir = data_dir()
+
+        return self._data_dir
+
+    @property
+    def bin_dir(self):  # type: () -> Path
+        if self._data_dir is not None:
+            return self._data_dir
+
+        from poetry.utils._compat import WINDOWS
+
+        if os.getenv("POETRY_HOME"):
+            return Path(os.getenv("POETRY_HOME"), "bin").expanduser()
+
+        user_base = site.getuserbase()
+
+        if WINDOWS:
+            bin_dir = os.path.join(user_base, "Scripts")
+        else:
+            bin_dir = os.path.join(user_base, "bin")
+
+        self._bin_dir = Path(bin_dir)
+
+        return self._bin_dir
+
+    @property
+    def pool(self):
+        if self._pool is not None:
+            return self._pool
+
+        from poetry.repositories.pool import Pool
+        from poetry.repositories.pypi_repository import PyPiRepository
+
+        pool = Pool()
+        pool.add_repository(PyPiRepository(fallback=False))
+
+        self._pool = pool
+
+        return self._pool
+
     def handle(self):
         from poetry.__version__ import __version__
         from poetry.core.semver import Version
-        from poetry.repositories.pypi_repository import PyPiRepository
+        from poetry.utils.env import EnvManager
+
+        new_update_method = False
+        try:
+            self._check_recommended_installation()
+        except RuntimeError as e:
+            env = EnvManager.get_system_env(naive=True)
+            try:
+                env.path.relative_to(self.data_dir)
+            except ValueError:
+                raise e
 
-        self._check_recommended_installation()
+            new_update_method = True
 
         version = self.argument("version")
         if not version:
             version = ">=" + __version__
 
-        repo = PyPiRepository(fallback=False)
+        repo = self.pool.repositories[0]
         packages = repo.find_packages(
             Dependency("poetry", version, allows_prereleases=self.option("preview"))
         )
@@ -127,6 +191,9 @@ def handle(self):
             self.line("You are using the latest version")
             return
 
+        if new_update_method:
+            return self.update_with_new_method(release.version)
+
         self.update(release)
 
     def update(self, release):
@@ -165,6 +232,18 @@ def update(self, release):
             )
         )
 
+    def update_with_new_method(self, version):
+        self.line("Updating <c1>Poetry</c1> to <c2>{}</c2>".format(version))
+        self.line("")
+
+        self._update_with_new_method(version)
+        self._make_bin()
+
+        self.line("")
+        self.line(
+            "<c1>Poetry</c1> (<c2>{}</c2>) is installed now. Great!".format(version)
+        )
+
     def _update(self, version):
         from poetry.utils.helpers import temporary_directory
 
@@ -235,6 +314,68 @@ def _update(self, version):
             finally:
                 gz.close()
 
+    def _update_with_new_method(self, version):
+        from poetry.config.config import Config
+        from poetry.core.packages.dependency import Dependency
+        from poetry.core.packages.project_package import ProjectPackage
+        from poetry.installation.installer import Installer
+        from poetry.packages.locker import NullLocker
+        from poetry.repositories.installed_repository import InstalledRepository
+        from poetry.utils.env import EnvManager
+
+        env = EnvManager.get_system_env()
+        installed = InstalledRepository.load(env)
+
+        root = ProjectPackage("poetry-updater", "0.0.0")
+        root.python_versions = ".".join(str(c) for c in env.version_info[:3])
+        root.add_dependency(Dependency("poetry", version.text))
+
+        installer = Installer(
+            self.io,
+            env,
+            root,
+            NullLocker(self.data_dir.joinpath("poetry.lock"), {}),
+            self.pool,
+            Config(),
+            installed=installed,
+        )
+        installer.update(True)
+        installer.run()
+
+    def _make_bin(self):
+        from poetry.utils._compat import WINDOWS
+
+        self.line("")
+        self.line("Updating the <c1>poetry</c1> script")
+
+        self.bin_dir.mkdir(parents=True, exist_ok=True)
+
+        script = "poetry"
+        target_script = "venv/bin/poetry"
+        if WINDOWS:
+            script = "poetry.exe"
+            target_script = "venv/Scripts/poetry.exe"
+
+        if self.bin_dir.joinpath(script).exists():
+            self.bin_dir.joinpath(script).unlink()
+
+        if not PY2 and not WINDOWS:
+            try:
+                self.bin_dir.joinpath(script).symlink_to(
+                    self.data_dir.joinpath(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)
+                )
+        else:
+            shutil.copy(
+                str(self.data_dir.joinpath(target_script)),
+                str(self.bin_dir.joinpath(script)),
+            )
+
     def process(self, *args):
         return subprocess.check_output(list(args), stderr=subprocess.STDOUT)
 

poetry/locations.py

@@ -1,9 +1,19 @@
+import os
+
 from .utils._compat import Path
 from .utils.appdirs import user_cache_dir
 from .utils.appdirs import user_config_dir
+from .utils.appdirs import user_data_dir
 
 
 CACHE_DIR = user_cache_dir("pypoetry")
 CONFIG_DIR = user_config_dir("pypoetry")
 
 REPOSITORY_CACHE_DIR = Path(CACHE_DIR) / "cache" / "repositories"
+
+
+def data_dir():  # type: () -> Path
+    if os.getenv("POETRY_HOME"):
+        return Path(os.getenv("POETRY_HOME")).expanduser()
+
+    return Path(user_data_dir("pypoetry", roaming=True))

poetry/packages/locker.py

@@ -586,3 +586,8 @@ def _dump_package(self, package):  # type: (Package) -> dict
                 data["develop"] = package.develop
 
         return data
+
+
+class NullLocker(Locker):
+    def set_lock_data(self, root, packages):  # type: (Package, List[Package]) -> None
+        pass

poetry/utils/env.py

@@ -472,7 +472,7 @@ def get(self, reload=False):  # type: (bool) -> Env
             create_venv = self._poetry.config.get("virtualenvs.create", True)
 
             if not create_venv:
-                return SystemEnv(Path(sys.prefix))
+                return self.get_system_env()
 
             venv_path = self._poetry.config.get("virtualenvs.path")
             if venv_path is None:
@@ -485,7 +485,7 @@ def get(self, reload=False):  # type: (bool) -> Env
             venv = venv_path / name
 
             if not venv.exists():
-                return SystemEnv(Path(sys.prefix))
+                return self.get_system_env()
 
             return VirtualEnv(venv)
 
@@ -790,7 +790,7 @@ def create_venv(
         p_venv = os.path.normcase(str(venv))
         if any(p.startswith(p_venv) for p in paths):
             # Running properly in the virtualenv, don't need to do anything
-            return SystemEnv(Path(sys.prefix), self.get_base_prefix())
+            return SystemEnv(Path(sys.prefix), Path(self.get_base_prefix()))
 
         return VirtualEnv(venv)
 
@@ -833,7 +833,33 @@ def remove_venv(cls, path):  # type: (Union[Path,str]) -> None
             elif file_path.is_dir():
                 shutil.rmtree(str(file_path))
 
-    def get_base_prefix(self):  # type: () -> Path
+    @classmethod
+    def get_system_env(cls, naive=False):  # type: (bool) -> "SystemEnv"
+        """
+        Retrieve the current Python environment.
+        This can be the base Python environment or an activated virtual environment.
+        This method also works around the issue that the virtual environment
+        used by Poetry internally (when installed via the custom installer)
+        is incorrectly detected as the system environment. Note that this workaround
+        happens only when `naive` is False since there are times where we actually
+        want to retrieve Poetry's custom virtual environment
+        (e.g. plugin installation or self update).
+        """
+        prefix, base_prefix = Path(sys.prefix), Path(cls.get_base_prefix())
+        if naive is False:
+            from poetry.locations import data_dir
+
+            try:
+                prefix.relative_to(data_dir())
+            except ValueError:
+                pass
+            else:
+                prefix = base_prefix
+
+        return SystemEnv(prefix)
+
+    @classmethod
+    def get_base_prefix(cls):  # type: () -> str
         if hasattr(sys, "real_prefix"):
             return sys.real_prefix
 
@@ -993,7 +1019,7 @@ def supported_tags(self):  # type: () -> List[Tag]
         return self._supported_tags
 
     @classmethod
-    def get_base_prefix(cls):  # type: () -> Path
+    def get_base_prefix(cls):  # type: () -> str
         if hasattr(sys, "real_prefix"):
             return sys.real_prefix