Add index url handling for upgrades and downgrades

StabilityMatrix.Avalonia/ViewModels/Dialogs/PythonPackagesItemViewModel.cs

@@ -60,6 +60,24 @@ public partial class PythonPackagesItemViewModel : ViewModelBase
         CanDowngrade = precedence < 0;
     }
 
+    /// <summary>
+    /// Return the known index URL for a package, currently this is torch, torchvision and torchaudio
+    /// </summary>
+    public static string? GetKnownIndexUrl(string packageName, string version)
+    {
+        var torchPackages = new[] { "torch", "torchvision", "torchaudio" };
+        if (torchPackages.Contains(packageName) && version.Contains('+'))
+        {
+            // Get the metadata for the current version (everything after the +)
+            var indexName = version.Split('+', 2).Last();
+
+            var indexUrl = $"https://download.pytorch.org/whl/{indexName}";
+            return indexUrl;
+        }
+
+        return null;
+    }
+
     /// <summary>
     /// Loads the pip show result if not already loaded
     /// </summary>
@@ -84,14 +102,8 @@ public partial class PythonPackagesItemViewModel : ViewModelBase
 
                 PipShowResult = await venvRunner.PipShow(Package.Name);
 
-                // Special case, include index for torch packages with + in the version
-                var torchPackages = new[] { "torch", "torchvision", "torchaudio" };
-                if (torchPackages.Contains(Package.Name) && Package.Version.Contains('+'))
-                {
-                    // Get the metadata for the current version (everything after the +)
-                    var indexName = Package.Version.Split('+', 2).Last();
-
-                    var indexUrl = $"https://download.pytorch.org/whl/{indexName}";
+                // Attempt to get known index url
+                var indexUrl = GetKnownIndexUrl(Package.Name, Package.Version);
 
                 if (await venvRunner.PipIndex(Package.Name, indexUrl) is { } pipIndexResult)
                 {
@@ -99,15 +111,6 @@ public partial class PythonPackagesItemViewModel : ViewModelBase
                     SelectedVersion = Package.Version;
                 }
             }
-                else
-                {
-                    if (await venvRunner.PipIndex(Package.Name) is { } pipIndexResult)
-                    {
-                        AvailableVersions = pipIndexResult.AvailableVersions;
-                        SelectedVersion = Package.Version;
-                    }
-                }
-            }
         }
         finally
         {

StabilityMatrix.Avalonia/ViewModels/Dialogs/PythonPackagesViewModel.cs

@@ -23,6 +23,7 @@ using StabilityMatrix.Core.Helper;
 using StabilityMatrix.Core.Models;
 using StabilityMatrix.Core.Models.FileInterfaces;
 using StabilityMatrix.Core.Models.PackageModification;
+using StabilityMatrix.Core.Processes;
 using StabilityMatrix.Core.Python;
 
 namespace StabilityMatrix.Avalonia.ViewModels.Dialogs;
@@ -136,13 +137,22 @@ public partial class PythonPackagesViewModel : ContentDialogViewModelBase
     private Task ModifySelectedPackage(PythonPackagesItemViewModel? item)
     {
         return item?.SelectedVersion != null
-            ? UpgradePackageVersion(item.Package.Name, item.SelectedVersion, item.CanDowngrade)
+            ? UpgradePackageVersion(
+                item.Package.Name,
+                item.SelectedVersion,
+                PythonPackagesItemViewModel.GetKnownIndexUrl(
+                    item.Package.Name,
+                    item.SelectedVersion
+                ),
+                isDowngrade: item.CanDowngrade
+            )
             : Task.CompletedTask;
     }
 
     private async Task UpgradePackageVersion(
         string packageName,
         string version,
+        string? extraIndexUrl = null,
         bool isDowngrade = false
     )
     {
@@ -169,13 +179,20 @@ public partial class PythonPackagesViewModel : ContentDialogViewModelBase
             return;
         }
 
+        var args = new ProcessArgsBuilder("install", $"{packageName}=={version}");
+
+        if (extraIndexUrl != null)
+        {
+            args = args.AddArg(("--extra-index-url", extraIndexUrl));
+        }
+
         var steps = new List<IPackageStep>
         {
             new PipStep
             {
                 VenvDirectory = VenvPath,
                 WorkingDirectory = VenvPath.Parent,
-                Args = new[] { "install", $"{packageName}=={version}" }
+                Args = args
             }
         };