11from __future__ import annotations
22
33import logging
4- import os
54from dataclasses import dataclass , replace
65from enum import Enum
76from typing import TYPE_CHECKING , Final
@@ -49,7 +48,7 @@ def upgrade( # ruff:ignore[too-many-arguments] # mirrors the CLI's flat upgrad
4948 results : Final [list [PackageUpgradeResult ]] = []
5049 install_messages : Final [list [OutputMessage ]] = []
5150
52- for venv_dir in venv_dirs .values ():
51+ for package_spec , venv_dir in venv_dirs .items ():
5352 with VenvContainer (venv_dir .parent ).venv_lock (venv_dir ) as venv_lock :
5453 results .extend (
5554 _upgrade_venv (
@@ -59,6 +58,7 @@ def upgrade( # ruff:ignore[too-many-arguments] # mirrors the CLI's flat upgrad
5958 include_injected = include_injected ,
6059 force = force ,
6160 install = install ,
61+ package_spec = package_spec ,
6262 venv_args = venv_args ,
6363 python = python ,
6464 python_flag_passed = python_flag_passed ,
@@ -208,6 +208,7 @@ def _upgrade_venv( # ruff:ignore[too-many-arguments] # upgrade forwards the fu
208208 include_injected : bool ,
209209 force : bool ,
210210 install : bool = False ,
211+ package_spec : str | None = None ,
211212 venv_args : list [str ] | None = None ,
212213 python : str | None = None ,
213214 python_flag_passed : bool = False ,
@@ -225,7 +226,9 @@ def _upgrade_venv( # ruff:ignore[too-many-arguments] # upgrade forwards the fu
225226 venv_dir = None ,
226227 venv_args = venv_args or [],
227228 package_names = None ,
228- package_specs = [str (venv_dir ).split (os .path .sep )[- 1 ]],
229+ # the venv directory only carries the package name; the original argument keeps
230+ # any version specifier or extras the user asked to install
231+ package_specs = [package_spec or venv_dir .name ],
229232 local_bin_dir = paths .ctx .bin_dir ,
230233 local_man_dir = paths .ctx .man_dir ,
231234 python = python ,
0 commit comments