Refactor helpers for Package.load (#4243)

diff --git a/lib/src/command/upgrade.dart b/lib/src/command/upgrade.dart
index 8d3c966..3c40b67 100644
--- a/lib/src/command/upgrade.dart
+++ b/lib/src/command/upgrade.dart
@@ -279,10 +279,6 @@
 
   Future<void> _runUpgradeMajorVersions() async {
     final toUpgrade = _directDependenciesToUpgrade();
-    final workspace = {
-      for (final package in entrypoint.workspaceRoot.transitiveWorkspace)
-        package.dir: package,
-    };
     // Solve [resolvablePubspec] in-memory and consolidate the resolved
     // versions of the packages into a map for quick searching.
     final resolvedPackages = <String, PackageId>{};
@@ -292,16 +288,8 @@
         return await resolveVersions(
           SolveType.upgrade,
           cache,
-          Package.load(
-            entrypoint.workspaceRoot.dir,
-            entrypoint.cache.sources,
-            withPubspecOverrides: true,
-            loadPubspec: (
-              path, {
-              expectedName,
-              required withPubspecOverrides,
-            }) =>
-                stripVersionBounds(workspace[path]!.pubspec),
+          entrypoint.workspaceRoot.transformWorkspace(
+            (package) => stripVersionBounds(package.pubspec),
           ),
         );
       },
@@ -354,15 +342,9 @@
       final solveResult = await resolveVersions(
         SolveType.upgrade,
         cache,
-        Package.load(
-          entrypoint.workspaceRoot.dir,
-          entrypoint.cache.sources,
-          loadPubspec: (path, {expectedName, required withPubspecOverrides}) {
-            final package = workspace[path]!;
-            final changesForPackage = changes[package] ?? {};
-            return applyChanges(package.pubspec, changesForPackage);
-          },
-        ),
+        entrypoint.workspaceRoot.transformWorkspace((package) {
+          return applyChanges(package.pubspec, changes[package] ?? {});
+        }),
       );
       changes = tighten(
         entrypoint,
@@ -388,7 +370,7 @@
         }
       }
     }
-    await entrypoint.withUpdatedPubspecs({
+    await entrypoint.withUpdatedRootPubspecs({
       for (final MapEntry(key: package, value: changesForPackage)
           in changes.entries)
         package: applyChanges(package.pubspec, changesForPackage),
diff --git a/lib/src/entrypoint.dart b/lib/src/entrypoint.dart
index 94f8dff..60515f4 100644
--- a/lib/src/entrypoint.dart
+++ b/lib/src/entrypoint.dart
@@ -105,7 +105,6 @@
       if (pubspec.resolution == Resolution.none) {
         root = Package.load(
           dir,
-          cache.sources,
           loadPubspec: (
             path, {
             expectedName,
@@ -281,8 +280,8 @@
       for (var packageEntry in packageConfig.nonInjectedPackages)
         packageEntry.name: Package.load(
           packageEntry.resolvedRootDir(packageConfigPath),
-          cache.sources,
           expectedName: packageEntry.name,
+          loadPubspec: Pubspec.loadRootWithSources(cache.sources),
         ),
     };
     packages[workspaceRoot.name] = workspaceRoot;
@@ -336,28 +335,10 @@
   }
 
   /// Creates an entrypoint at the same location, but with each pubspec in
-  /// [updatedPubspec] replacing the with one for the corresponding package.
-  Entrypoint withUpdatedPubspecs(Map<Package, Pubspec> updatedPubspecs) {
-    final existingPubspecs = <String, Pubspec>{};
-    // First extract all pubspecs from the workspace.
-    for (final package in workspaceRoot.transitiveWorkspace) {
-      existingPubspecs[package.dir] =
-          updatedPubspecs[package] ?? package.pubspec;
-    }
-    final newWorkspaceRoot = Package.load(
-      workspaceRoot.dir,
-      cache.sources,
-      loadPubspec: (
-        dir, {
-        expectedName,
-        required withPubspecOverrides,
-      }) =>
-          existingPubspecs[dir] ??
-          Pubspec.load(
-            dir,
-            cache.sources,
-            containingDescription: RootDescription(dir),
-          ),
+  /// [updatedPubspecs] replacing the with one for the corresponding package.
+  Entrypoint withUpdatedRootPubspecs(Map<Package, Pubspec> updatedPubspecs) {
+    final newWorkspaceRoot = workspaceRoot.transformWorkspace(
+      (package) => updatedPubspecs[package] ?? package.pubspec,
     );
     final newWorkPackage = newWorkspaceRoot.transitiveWorkspace
         .firstWhere((package) => package.dir == workPackage.dir);
@@ -375,7 +356,7 @@
   /// Creates an entrypoint at the same location, that will use [pubspec] for
   /// resolution of the [workPackage].
   Entrypoint withWorkPubspec(Pubspec pubspec) {
-    return withUpdatedPubspecs({workPackage: pubspec});
+    return withUpdatedRootPubspecs({workPackage: pubspec});
   }
 
   /// Creates an entrypoint given package and lockfile objects.
@@ -1108,7 +1089,10 @@
       }
       var touchedLockFile = false;
       late final lockFile = _loadLockFile(lockFilePath, cache);
-      late final root = Package.load(dir, cache.sources);
+      late final root = Package.load(
+        dir,
+        loadPubspec: Pubspec.loadRootWithSources(cache.sources),
+      );
 
       if (!lockfileNewerThanPubspecs) {
         if (isLockFileUpToDate(lockFile, root)) {
diff --git a/lib/src/package.dart b/lib/src/package.dart
index dc21f8c..ff6291a 100644
--- a/lib/src/package.dart
+++ b/lib/src/package.dart
@@ -14,8 +14,6 @@
 import 'log.dart' as log;
 import 'package_name.dart';
 import 'pubspec.dart';
-import 'source/root.dart';
-import 'system_cache.dart';
 import 'utils.dart';
 
 /// A Package is a [Pubspec] and a directory where it belongs that can be used
@@ -151,22 +149,15 @@
   /// used to override a pubspec in memory for trying out an alternative
   /// resolution.
   factory Package.load(
-    String dir,
-    SourceRegistry sources, {
+    String dir, {
     bool withPubspecOverrides = false,
     String? expectedName,
-    Pubspec Function(
+    required Pubspec Function(
       String path, {
       String? expectedName,
       required bool withPubspecOverrides,
-    })? loadPubspec,
+    }) loadPubspec,
   }) {
-    loadPubspec ??=
-        (path, {expectedName, required withPubspecOverrides}) => Pubspec.load(
-              path,
-              sources,
-              containingDescription: RootDescription(path),
-            );
     final pubspec = loadPubspec(
       dir,
       withPubspecOverrides: withPubspecOverrides,
@@ -178,7 +169,6 @@
         try {
           return Package.load(
             p.join(dir, workspacePath),
-            sources,
             loadPubspec: loadPubspec,
             withPubspecOverrides: withPubspecOverrides,
           );
@@ -365,6 +355,26 @@
       isDir: (dir) => dirExists(resolve(dir)),
     ).map(resolve).toList();
   }
+
+  /// Applies [transform] to each package in the workspace and returns a derived
+  /// package.
+  Package transformWorkspace(
+    Pubspec Function(Package) transform,
+  ) {
+    final workspace = {
+      for (final package in transitiveWorkspace) package.dir: package,
+    };
+    return Package.load(
+      dir,
+      withPubspecOverrides: true,
+      loadPubspec: (
+        path, {
+        expectedName,
+        required withPubspecOverrides,
+      }) =>
+          transform(workspace[path]!),
+    );
+  }
 }
 
 /// Reports an error if the graph of the workspace rooted at [root] is not a
diff --git a/lib/src/pubspec.dart b/lib/src/pubspec.dart
index 2a951ad..deabc90 100644
--- a/lib/src/pubspec.dart
+++ b/lib/src/pubspec.dart
@@ -58,7 +58,7 @@
   /// [devDependencies].
   ///
   /// This will be null if this was created using [Pubspec] or [Pubspec.empty].
-  final SourceRegistry _sources;
+  final SourceRegistry sources;
 
   /// It is used to resolve relative paths. And to resolve path-descriptions
   /// from a git dependency as git-descriptions.
@@ -121,7 +121,7 @@
       _dependencies ??= _parseDependencies(
         'dependencies',
         fields.nodes['dependencies'],
-        _sources,
+        sources,
         languageVersion,
         _packageName,
         _containingDescription,
@@ -134,7 +134,7 @@
       _devDependencies ??= _parseDependencies(
         'dev_dependencies',
         fields.nodes['dev_dependencies'],
-        _sources,
+        sources,
         languageVersion,
         _packageName,
         _containingDescription,
@@ -166,7 +166,7 @@
         _dependencyOverrides = _parseDependencies(
           'dependency_overrides',
           pubspecOverridesFields.nodes['dependency_overrides'],
-          _sources,
+          sources,
           languageVersion,
           _packageName,
           _containingDescription,
@@ -177,7 +177,7 @@
     return _dependencyOverrides ??= _parseDependencies(
       'dependency_overrides',
       fields.nodes['dependency_overrides'],
-      _sources,
+      sources,
       languageVersion,
       _packageName,
       _containingDescription,
@@ -298,6 +298,26 @@
     );
   }
 
+  /// Convenience helper to pass to [Package.load].
+  static Pubspec Function(
+    String dir, {
+    String? expectedName,
+    required bool withPubspecOverrides,
+  }) loadRootWithSources(SourceRegistry sources) {
+    return (
+      String dir, {
+      String? expectedName,
+      required bool withPubspecOverrides,
+    }) =>
+        Pubspec.load(
+          dir,
+          sources,
+          expectedName: expectedName,
+          allowOverridesFile: withPubspecOverrides,
+          containingDescription: RootDescription(dir),
+        );
+  }
+
   Pubspec(
     String name, {
     Version? version,
@@ -322,7 +342,7 @@
         _givenSdkConstraints = sdkConstraints ??
             UnmodifiableMapView({'dart': SdkConstraint(VersionConstraint.any)}),
         _includeDefaultSdkConstraint = false,
-        _sources = sources ??
+        sources = sources ??
             ((String? name) => throw StateError('No source registry given')),
         _overridesFileFields = null,
         // This is a dummy value.
@@ -343,7 +363,7 @@
   /// [location] is the location from which this pubspec was loaded.
   Pubspec.fromMap(
     Map fields,
-    this._sources, {
+    this.sources, {
     YamlMap? overridesFields,
     String? expectedName,
     Uri? location,
diff --git a/lib/src/source/git.dart b/lib/src/source/git.dart
index 7f863aa..783dc48 100644
--- a/lib/src/source/git.dart
+++ b/lib/src/source/git.dart
@@ -443,7 +443,10 @@
 
             var packageDir = p.join(revisionCachePath, relative);
             try {
-              return Package.load(packageDir, cache.sources);
+              return Package.load(
+                packageDir,
+                loadPubspec: Pubspec.loadRootWithSources(cache.sources),
+              );
             } catch (error, stackTrace) {
               log.error('Failed to load package', error, stackTrace);
               var name = p.basename(revisionCachePath).split('-').first;
diff --git a/lib/src/source/hosted.dart b/lib/src/source/hosted.dart
index b51baa9..f3def3b 100644
--- a/lib/src/source/hosted.dart
+++ b/lib/src/source/hosted.dart
@@ -1277,7 +1277,12 @@
         var packages = <Package>[];
         for (var entry in listDir(serverDir)) {
           try {
-            packages.add(Package.load(entry, cache.sources));
+            packages.add(
+              Package.load(
+                entry,
+                loadPubspec: Pubspec.loadRootWithSources(cache.sources),
+              ),
+            );
           } catch (error, stackTrace) {
             log.error('Failed to load package', error, stackTrace);
             final id = _idForBasename(
@@ -1384,7 +1389,10 @@
         .where(_looksLikePackageDir)
         .map((entry) {
           try {
-            return Package.load(entry, cache.sources);
+            return Package.load(
+              entry,
+              loadPubspec: Pubspec.loadRootWithSources(cache.sources),
+            );
           } catch (error, stackTrace) {
             log.fine('Failed to load package from $entry:\n'
                 '$error\n'
diff --git a/lib/src/system_cache.dart b/lib/src/system_cache.dart
index e20a60b..b98bbb9 100644
--- a/lib/src/system_cache.dart
+++ b/lib/src/system_cache.dart
@@ -111,20 +111,11 @@
   ///
   /// Throws an [ArgumentError] if [id] has an invalid source.
   Package load(PackageId id) {
-    return Package.load(getDirectory(id), sources, expectedName: id.name);
-  }
-
-  Package loadCached(PackageId id) {
-    final source = id.description.description.source;
-    if (source is CachedSource) {
-      return Package.load(
-        source.getDirectoryInCache(id, this),
-        sources,
-        expectedName: id.name,
-      );
-    } else {
-      throw ArgumentError('Call only on Cached ids.');
-    }
+    return Package.load(
+      getDirectory(id),
+      loadPubspec: Pubspec.loadRootWithSources(sources),
+      expectedName: id.name,
+    );
   }
 
   /// Create a new temporary directory within the system cache.
diff --git a/test/ascii_tree_test.dart b/test/ascii_tree_test.dart
index 115cbc6..7c1eb9b 100644
--- a/test/ascii_tree_test.dart
+++ b/test/ascii_tree_test.dart
@@ -4,6 +4,7 @@
 
 import 'package:pub/src/ascii_tree.dart' as tree;
 import 'package:pub/src/package.dart';
+import 'package:pub/src/pubspec.dart';
 import 'package:pub/src/utils.dart';
 import 'package:test/test.dart';
 
@@ -61,9 +62,11 @@
         file('path.dart', bytes(100)),
       ]),
     ]).create();
-    var files =
-        Package.load(path(appPath), (name) => throw UnimplementedError())
-            .listFiles();
+    var files = Package.load(
+      path(appPath),
+      loadPubspec:
+          Pubspec.loadRootWithSources((name) => throw UnimplementedError()),
+    ).listFiles();
     ctx.expectNextSection(
       tree.fromFiles(files, baseDir: path(appPath), showFileSizes: true),
     );