From 08b8f378ce82890e35b29566293e8f31f1e34a44 Mon Sep 17 00:00:00 2001
From: chrysle <i.kant@posteo.de>
Date: Sun, 25 Feb 2024 17:58:14 +0100
Subject: [PATCH] Add `--install` option to `pipx upgrade` command

---
 src/pipx/commands/upgrade.py | 49 ++++++++++++++++++++++++++++--------
 src/pipx/main.py             |  8 ++++++
 tests/test_upgrade.py        |  6 +++++
 3 files changed, 52 insertions(+), 11 deletions(-)

diff --git a/src/pipx/commands/upgrade.py b/src/pipx/commands/upgrade.py
index 1530a30cdf..ccb0ebd1ae 100644
--- a/src/pipx/commands/upgrade.py
+++ b/src/pipx/commands/upgrade.py
@@ -1,8 +1,8 @@
 import logging
 from pathlib import Path
-from typing import List, Sequence
+from typing import List, Optional, Sequence
 
-from pipx import constants
+from pipx import commands, constants
 from pipx.colors import bold, red
 from pipx.commands.common import expose_resources_globally
 from pipx.constants import EXIT_CODE_OK, ExitCode
@@ -101,15 +101,38 @@ def _upgrade_venv(
     include_injected: bool,
     upgrading_all: bool,
     force: bool,
+    install: bool = False,
+    python: Optional[str] = None,
 ) -> int:
-    """Returns number of packages with changed versions."""
+    """Return number of packages with changed versions."""
     if not venv_dir.is_dir():
-        raise PipxError(
-            f"""
-            Package is not installed. Expected to find {str(venv_dir)}, but it
-            does not exist.
-            """
-        )
+        if install:
+            commands.install(
+                venv_dir=None,
+                venv_args=[],
+                package_names=None,
+                package_specs=[str(venv_dir).split("/")[-1]],
+                local_bin_dir=constants.LOCAL_BIN_DIR,
+                local_man_dir=constants.LOCAL_MAN_DIR,
+                python=python,
+                pip_args=pip_args,
+                verbose=verbose,
+                force=force,
+                reinstall=False,
+                include_dependencies=False,
+                preinstall_packages=None,
+            )
+            return 0
+        else:
+            raise PipxError(
+                f"""
+                Package is not installed. Expected to find {str(venv_dir)}, but it
+                does not exist.
+                """
+            )
+
+    if python and not install:
+        logger.info("Ignoring --python as not combined with --install")
 
     venv = Venv(venv_dir, verbose=verbose)
 
@@ -154,13 +177,15 @@ def _upgrade_venv(
 
 def upgrade(
     venv_dir: Path,
+    python: Optional[str],
     pip_args: List[str],
     verbose: bool,
     *,
     include_injected: bool,
     force: bool,
+    install: bool,
 ) -> ExitCode:
-    """Returns pipx exit code."""
+    """Return pipx exit code."""
 
     _ = _upgrade_venv(
         venv_dir,
@@ -169,6 +194,8 @@ def upgrade(
         include_injected=include_injected,
         upgrading_all=False,
         force=force,
+        install=install,
+        python=python,
     )
 
     # Any error in upgrade will raise PipxError (e.g. from venv.upgrade_package())
@@ -194,7 +221,7 @@ def upgrade_all(
             venvs_upgraded += _upgrade_venv(
                 venv_dir,
                 venv.pipx_metadata.main_package.pip_args,
-                verbose,
+                verbose=verbose,
                 include_injected=include_injected,
                 upgrading_all=True,
                 force=force,
diff --git a/src/pipx/main.py b/src/pipx/main.py
index cfea91208f..93471ff196 100644
--- a/src/pipx/main.py
+++ b/src/pipx/main.py
@@ -267,10 +267,12 @@ def run_pipx_command(args: argparse.Namespace) -> ExitCode:  # noqa: C901
     elif args.command == "upgrade":
         return commands.upgrade(
             venv_dir,
+            args.python,
             pip_args,
             verbose,
             include_injected=args.include_injected,
             force=args.force,
+            install=args.install,
         )
     elif args.command == "upgrade-all":
         return commands.upgrade_all(
@@ -486,6 +488,12 @@ def _add_upgrade(subparsers, venv_completer: VenvCompleter, shared_parser: argpa
         help="Modify existing virtual environment and files in PIPX_BIN_DIR and PIPX_MAN_DIR",
     )
     add_pip_venv_args(p)
+    p.add_argument(
+        "--install",
+        action="store_true",
+        help="Install package spec if missing",
+    )
+    add_python_options(p)
 
 
 def _add_upgrade_all(subparsers: argparse._SubParsersAction, shared_parser: argparse.ArgumentParser) -> None:
diff --git a/tests/test_upgrade.py b/tests/test_upgrade.py
index 816b7c4ed6..c89629da50 100644
--- a/tests/test_upgrade.py
+++ b/tests/test_upgrade.py
@@ -73,3 +73,9 @@ def test_upgrade_no_include_injected(pipx_temp_env, capsys):
     captured = capsys.readouterr()
     assert "upgraded package pylint" in captured.out
     assert "upgraded package black" not in captured.out
+
+
+def test_upgrade_install_missing(pipx_temp_env, capsys):
+    assert not run_pipx_cli(["upgrade", "pycowsay", "--install"])
+    captured = capsys.readouterr()
+    assert "installed package pycowsay" in captured.out