Ensure the self update command is compatible with the new installer

Author: sdispater

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,7 @@
 from cleo import option
 
 from poetry.core.packages import Dependency
+from poetry.utils._compat import Path
 
 from ..command import Command
 
@@ -60,6 +62,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 +84,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
 
-        self._check_recommended_installation()
+        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
+
+            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 +190,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 +231,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 +313,62 @@ 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()
+
+        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)
+            )
+
     def process(self, *args):
         return subprocess.check_output(list(args), stderr=subprocess.STDOUT)
 

poetry/packages/locker.py

@@ -348,7 +348,9 @@ def get_project_dependency_packages(
         if extra_package_names is not None:
             extra_package_names = set(
                 get_extra_package_names(
-                    repository.packages, self.lock_data.get("extras", {}), extras or (),
+                    repository.packages,
+                    self.lock_data.get("extras", {}),
+                    extras or (),
                 )
             )
 
@@ -586,3 +588,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

tests/console/commands/self/test_update.py

@@ -5,8 +5,13 @@
 from poetry.__version__ import __version__
 from poetry.core.packages.package import Package
 from poetry.core.semver.version import Version
+from poetry.factory import Factory
+from poetry.repositories.installed_repository import InstalledRepository
+from poetry.repositories.pool import Pool
+from poetry.repositories.repository import Repository
 from poetry.utils._compat import WINDOWS
 from poetry.utils._compat import Path
+from poetry.utils.env import EnvManager
 
 
 FIXTURES = Path(__file__).parent.joinpath("fixtures")
@@ -25,10 +30,13 @@ def test_self_update_should_install_all_necessary_elements(
     command = tester._command
 
     version = Version.parse(__version__).next_minor.text
-    mocker.patch(
-        "poetry.repositories.pypi_repository.PyPiRepository.find_packages",
-        return_value=[Package("poetry", version)],
-    )
+    repository = Repository()
+    repository.add_package(Package("poetry", version))
+
+    pool = Pool()
+    pool.add_repository(repository)
+
+    command._pool = pool
     mocker.patch.object(command, "_check_recommended_installation", return_value=None)
     mocker.patch.object(
         command, "_get_release_name", return_value="poetry-{}-darwin".format(version)
@@ -89,3 +97,57 @@ def test_self_update_should_install_all_necessary_elements(
 
     assert lib.exists()
     assert lib.joinpath("poetry").exists()
+
+
+def test_self_update_can_update_from_recommended_installation(
+    tester, http, mocker, environ, tmp_venv
+):
+    mocker.patch.object(EnvManager, "get_system_env", return_value=tmp_venv)
+
+    command = tester._command
+    command._data_dir = tmp_venv.path.parent
+
+    new_version = Version.parse(__version__).next_minor.text
+
+    old_poetry = Package("poetry", __version__)
+    old_poetry.add_dependency(Factory.create_dependency("cleo", "^0.8.2"))
+
+    new_poetry = Package("poetry", new_version)
+    new_poetry.add_dependency(Factory.create_dependency("cleo", "^1.0.0"))
+
+    installed_repository = Repository()
+    installed_repository.add_package(old_poetry)
+    installed_repository.add_package(Package("cleo", "0.8.2"))
+
+    repository = Repository()
+    repository.add_package(new_poetry)
+    repository.add_package(Package("cleo", "1.0.0"))
+
+    pool = Pool()
+    pool.add_repository(repository)
+
+    command._pool = pool
+
+    mocker.patch.object(InstalledRepository, "load", return_value=installed_repository)
+
+    tester.execute()
+
+    expected_output = """\
+Updating Poetry to {}
+
+Updating dependencies
+Resolving dependencies...
+
+Package operations: 0 installs, 2 updates, 0 removals
+
+  - Updating cleo (0.8.2 -> 1.0.0)
+  - Updating poetry ({} -> {})
+
+Updating the poetry script
+
+Poetry (1.2.0) is installed now. Great!
+""".format(
+        new_version, __version__, new_version
+    )
+
+    assert tester.io.fetch_output() == expected_output