Backport: install script improvments (#841)

Co-authored-by: Wizard1209 <[email protected]>

Author: Wizard1209

src/dipdup/cli.py

@@ -575,7 +575,7 @@ async def install(
     """Install DipDup for the current user."""
     import dipdup.install
 
-    dipdup.install.install(quiet, force, ref, path)
+    dipdup.install.install(quiet, force, None, ref, path)
 
 
 @cli.command()
@@ -605,4 +605,4 @@ async def update(
     """Update DipDup for the current user."""
     import dipdup.install
 
-    dipdup.install.install(quiet, force, None, None)
+    dipdup.install.install(quiet, force, None, None, None)

src/dipdup/install.py

@@ -14,22 +14,37 @@
 from shutil import rmtree
 from shutil import which
 from typing import Any
-from typing import Dict
 from typing import NoReturn
-from typing import Optional
-from typing import Set
+from typing import cast
 
 GITHUB = 'https://github.com/dipdup-io/dipdup.git'
 WHICH_CMDS = (
+    'python3.10',
     'python3',
     'pipx',
     'dipdup',
     'datamodel-codegen',
     'poetry',
+    'pdm',
     'pyvenv',
     'pyenv',
 )
 
+WELCOME_ASCII = """\0
+        ____   _         ____              
+       / __ \ (_)____   / __ \ __  __ ____ 
+      / / / // // __ \ / / / // / / // __ \\
+     / /_/ // // /_/ // /_/ // /_/ // /_/ /
+    /_____//_// .___//_____/ \__,_// .___/ 
+             /_/                  /_/      
+"""
+EPILOG = """\0
+Documentation:         https://dipdup.io/docs
+GitHub:                https://github.com/dipdup-io/dipdup
+Discord:               https://discord.gg/aG8XKuwsQd
+"""
+DIPDUP_LTS_VERSION = '6.5.12'
+
 
 class Colors:
     """ANSI color codes"""
@@ -55,16 +70,16 @@ def done(msg: str) -> NoReturn:
     sys.exit(0)
 
 
-def ask(msg: str, default: bool, quiet: bool) -> bool:
-    msg += ' [Y/n]' if default else ' [y/N]'
-    echo(msg, Colors.YELLOW)
-
-    if quiet:
-        return default
-    if default:
-        return input().lower() not in ('n', 'no')
-    else:
-        return input().lower() in ('y', 'yes')
+def ask(question: str, default: bool) -> bool:
+    """Ask user a yes/no question"""
+    while True:
+        answer = input(question + (' [Y/n] ' if default else ' [y/N] ')).lower().strip()
+        if not answer:
+            return default
+        if answer in ('n', 'no'):
+            return False
+        if answer in ('y', 'yes'):
+            return True
 
 
 # NOTE: DipDup has `tabulate` dep, don't use this one elsewhere
@@ -73,42 +88,40 @@ def _tab(text: str, indent: int = 20) -> str:
 
 
 class DipDupEnvironment:
-    def __init__(self, quiet: bool = False) -> None:
+    def __init__(self) -> None:
         self._os = os.uname().sysname
         self._arch = os.uname().machine
-        self._quiet = quiet
-        self._commands: Dict[str, Optional[str]] = {}
-        self._pipx_packages: Set[str] = set()
+        self._commands: dict[str, str | None] = {}
+        self._pipx_packages: set[str] = set()
 
     def refresh(self) -> None:
-        if not self._quiet and not self._commands:
-            print()
-            print(_tab('OS:') + self._os)
-            print(_tab('Arch:') + self._arch)
-            print(_tab('Python:') + sys.version)
-            print(_tab('PATH:') + os.environ['PATH'])
-            print()
-
         for command in WHICH_CMDS:
             old, new = self._commands.get(command), which(command)
             if old == new:
                 continue
             self._commands[command] = new
-            self._quiet or print(_tab(command) + (new or ''))
 
+    def print(self) -> None:
+        print()
+        print(WELCOME_ASCII)
+        print(EPILOG)
         print()
+        print(_tab('OS:') + f'{self._os} ({self._arch})')
+        print(_tab('Python:') + sys.version)
+        print(_tab('PATH:') + os.environ['PATH'])
+        print(_tab('PYTHONPATH:') + os.environ.get('PYTHONPATH', ''))
+        print()
+        for command, path in self._commands.items():
+            print(_tab(f'{command}:') + (path or ''))
+        print(_tab('pipx packages:') + ', '.join(self._pipx_packages) + '\n')
 
     def refresh_pipx(self) -> None:
         """Get installed pipx packages"""
         self.ensure_pipx()
         pipx_packages_raw = self.run_cmd('pipx', 'list', '--short', capture_output=True).stdout
         self._pipx_packages = {p.split()[0].decode() for p in pipx_packages_raw.splitlines()}
-        self._quiet or print(_tab('pipx packages:') + ', '.join(self._pipx_packages) + '\n')
-
-    def check(self) -> None:
-        if not sys.version.startswith('3.10'):
-            fail('DipDup requires Python 3.10')
 
+    def prepare(self) -> None:
         # NOTE: Show warning if user is root
         if os.geteuid() == 0:
             echo('WARNING: Running as root, this is not generally recommended', Colors.YELLOW)
@@ -120,109 +133,136 @@ def check(self) -> None:
         self.refresh()
         self.refresh_pipx()
 
-        if self._commands.get('pyenv'):
-            echo('WARNING: pyenv is installed, this may cause issues', Colors.YELLOW)
-
     def run_cmd(self, cmd: str, *args: Any, **kwargs: Any) -> subprocess.CompletedProcess[bytes]:
         """Run command safely (relatively lol)"""
         if (found_cmd := self._commands.get(cmd)) is None:
             fail(f'Command not found: {cmd}')
-        args = (found_cmd,) + tuple(a for a in args if a)
+        args = (found_cmd, *tuple(a for a in args if a))
         try:
             return subprocess.run(
                 args,
                 **kwargs,
                 check=True,
             )
         except subprocess.CalledProcessError as e:
-            self._quiet or fail(f'{cmd} failed: {e.cmd} {e.returncode}')
-            raise
+            fail(f'{cmd} failed: {e.cmd} {e.returncode}')
 
     def ensure_pipx(self) -> None:
+        if not sys.version.startswith('3.10'):
+            fail('DipDup 6 requires Python 3.10')
+
         """Ensure pipx is installed for current user"""
         if self._commands.get('pipx'):
             return
 
-        if sys.prefix != sys.base_prefix:
-            fail("pipx can't be installed in virtualenv, run `deactivate` and try again")
-
         echo('Installing pipx')
-        self.run_cmd('python3', '-m', 'pip', 'install', '--user', '-q', 'pipx')
-        self.run_cmd('python3', '-m', 'pipx', 'ensurepath')
-        os.environ['PATH'] = os.environ['PATH'] + ':' + str(Path.home() / '.local' / 'bin')
-        os.execv(sys.executable, [sys.executable] + sys.argv)
+        if sys.base_prefix != sys.prefix:
+            self.run_cmd('python3.10', '-m', 'pip', 'install', '-q', 'pipx')
+        else:
+            self.run_cmd('python3.10', '-m', 'pip', 'install', '--user', '-q', 'pipx')
+        self.run_cmd('python3.10', '-m', 'pipx', 'ensurepath')
+        pipx_path = str(Path.home() / '.local' / 'bin')
+        os.environ['PATH'] = pipx_path + os.pathsep + os.environ['PATH']
+        self._commands['pipx'] = which('pipx')
 
 
 def install(
     quiet: bool,
     force: bool,
+    version: str | None,
     ref: str | None,
     path: str | None,
 ) -> None:
     """Install DipDup and its dependencies with pipx"""
     if ref and path:
         fail('Specify either ref or path, not both')
+    pkg = 'dipdup' if not version else f'dipdup=={version}'
+
+    if not any((version, ref, path)):
+        version = DIPDUP_LTS_VERSION
 
     env = DipDupEnvironment()
-    env.check()
+    env.prepare()
+    if not quiet:
+        env.print()
 
     force_str = '--force' if force else ''
     pipx_packages = env._pipx_packages
     pipx_dipdup = 'dipdup' in pipx_packages
     pipx_datamodel_codegen = 'datamodel-code-generator' in pipx_packages
     pipx_poetry = 'poetry' in pipx_packages
 
-    if pipx_dipdup:
-        echo('Updating DipDup')
-        env.run_cmd('pipx', 'upgrade', 'dipdup', force_str)
+    python_inter_pipx = cast(str, which('python3.10'))
+    if 'pyenv' in python_inter_pipx:
+        python_inter_pipx = (
+            subprocess.run(
+                ['pyenv', 'which', 'python3.10'],
+                capture_output=True,
+                text=True,
+            )
+            .stdout.strip()
+            .split('\n')[0]
+        )
+
+    if pipx_dipdup and not force:
+        if quiet or ask(f'Reinstall dipdup to {version}', False):
+            env.run_cmd('pipx', 'uninstall', 'dipdup')
+            env.run_cmd('pipx', 'install', '--python', python_inter_pipx, pkg, force_str)
     else:
         if path:
             echo(f'Installing DipDup from `{path}`')
-            env.run_cmd('pipx', 'install', path, force_str)
+            env.run_cmd('pipx', 'install', '--python', python_inter_pipx, path, force_str)
         elif ref:
-            echo(f'Installing DipDup from `{ref}`')
-            env.run_cmd('pipx', 'install', f'git+{GITHUB}@{ref}', force_str)
+            url = f'git+{GITHUB}@{ref}'
+            echo(f'Installing DipDup from `{url}`')
+            env.run_cmd('pipx', 'install', '--python', python_inter_pipx, url, force_str)
         else:
             echo('Installing DipDup from PyPI')
-            env.run_cmd('pipx', 'install', 'dipdup', force_str)
+            env.run_cmd('pipx', 'install', '--python', python_inter_pipx, pkg, force_str)
 
     if pipx_datamodel_codegen:
         env.run_cmd('pipx', 'upgrade', 'datamodel-code-generator', force_str)
     else:
-        env.run_cmd('pipx', 'install', 'datamodel-code-generator', force_str)
+        env.run_cmd('pipx', 'install', '--python', python_inter_pipx, 'datamodel-code-generator', force_str)
 
     if (legacy_poetry := Path(Path.home(), '.poetry')).exists():
         rmtree(legacy_poetry, ignore_errors=True)
         env.run_cmd('pipx', 'install', 'poetry', force_str)
     elif pipx_poetry:
         echo('Updating Poetry')
         env.run_cmd('pipx', 'upgrade', 'poetry', force_str)
-    elif ask('Install poetry? Optional for `dipdup new` command', True, quiet):
+    elif quiet or ask('Install poetry? Optional for `dipdup new` command', True):
         echo('Installing poetry')
         env.run_cmd('pipx', 'install', 'poetry', force_str)
         env._commands['poetry'] = which('poetry')
         pipx_poetry = True
 
     done(
-        'Done! DipDup is ready to use.\nRun `dipdup new` to create a new project or `dipdup` to see all available commands.'
+        'Done! DipDup is ready to use.\nRun `dipdup new` to create a new project or `dipdup` to see all available'
+        ' commands.'
     )
 
 
 def uninstall(quiet: bool) -> NoReturn:
     """Uninstall DipDup and its dependencies with pipx"""
     env = DipDupEnvironment()
-    env.check()
-
-    pipx_packages = env._pipx_packages
-
-    if 'dipdup' in pipx_packages:
-        echo('Uninstalling DipDup')
-        env.run_cmd('pipx', 'uninstall', 'dipdup')
+    env.prepare()
+    if not quiet:
+        env.print()
+
+    packages = (
+        ('dipdup', True),
+        ('datamodel-code-generator', False),
+        ('poetry', False),
+    )
+    for package, default in packages:
+        if package not in env._pipx_packages:
+            continue
+        if not quiet and not ask(f'Uninstall {package}?', default):
+            continue
 
-    if 'datamodel-code-generator' in pipx_packages:
-        if ask('Uninstall datamodel-code-generator?', True, quiet):
-            echo('Uninstalling datamodel-code-generator')
-            env.run_cmd('pipx', 'uninstall', 'datamodel-code-generator')
+        echo(f'Uninstalling {package}')
+        env.run_cmd('pipx', 'uninstall', package)
 
     done('Done! DipDup is uninstalled.')
 
@@ -233,17 +273,22 @@ def cli() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument('-q', '--quiet', action='store_true', help='Use default answers for all questions')
     parser.add_argument('-f', '--force', action='store_true', help='Force reinstall')
+    parser.add_argument('-v', '--version', help='Install DipDup from a specific version')
     parser.add_argument('-r', '--ref', help='Install DipDup from a specific git ref')
     parser.add_argument('-p', '--path', help='Install DipDup from a local path')
     parser.add_argument('-u', '--uninstall', action='store_true', help='Uninstall DipDup')
     args = parser.parse_args()
 
+    if not args.quiet:
+        sys.stdin = open('/dev/tty')  # noqa: PTH123
+
     if args.uninstall:
         uninstall(args.quiet)
     else:
         install(
             quiet=args.quiet,
             force=args.force,
+            version=args.version.strip() if args.version else None,
             ref=args.ref.strip() if args.ref else None,
             path=args.path.strip() if args.path else None,
         )