Model Python version as a PubGrub package (#745)

## Summary

This PR modifies the resolver to treat the Python version as a package,
which allows for better error messages (since we no longer treat
incompatible packages as if they "don't exist at all").

There are a few tricky pieces here...

First, we need to track both the interpreter's Python version and the
_target_ Python version, because we support resolving for other versions
via `--python 3.7`.

Second, we allow using incompatible wheels during resolution, as long as
there's a compatible source distribution. So we still need to test for
`requires-python` compatibility when selecting distributions.

This could use more testing, but it feels like an area where `packse`
would be more productive than writing PyPI tests.

Closes https://github.com/astral-sh/puffin/issues/406.

crates/puffin-cli/src/commands/pip_compile.rs

@@ -207,8 +207,16 @@ pub(crate) async fn pip_compile(
     );
 
     // Resolve the dependencies.
-    let resolver = Resolver::new(manifest, options, &markers, tags, &client, &build_dispatch)
-        .with_reporter(ResolverReporter::from(printer));
+    let resolver = Resolver::new(
+        manifest,
+        options,
+        &markers,
+        &interpreter,
+        tags,
+        &client,
+        &build_dispatch,
+    )
+    .with_reporter(ResolverReporter::from(printer));
     let resolution = match resolver.resolve().await {
         Err(puffin_resolver::ResolveError::NoSolution(err)) => {
             #[allow(clippy::print_stderr)]

crates/puffin-cli/src/commands/pip_install.rs

@@ -21,7 +21,7 @@ use puffin_dispatch::BuildDispatch;
 use puffin_installer::{
     BuiltEditable, Downloader, InstallPlan, Reinstall, ResolvedEditable, SitePackages,
 };
-use puffin_interpreter::Virtualenv;
+use puffin_interpreter::{Interpreter, Virtualenv};
 use puffin_normalize::PackageName;
 use puffin_resolver::{
     Manifest, PreReleaseMode, ResolutionGraph, ResolutionMode, ResolutionOptions, Resolver,
@@ -173,6 +173,7 @@ pub(crate) async fn pip_install(
         &editables,
         &site_packages,
         reinstall,
+        &interpreter,
         tags,
         markers,
         &client,
@@ -320,6 +321,7 @@ async fn resolve(
     editables: &[BuiltEditable],
     site_packages: &SitePackages<'_>,
     reinstall: &Reinstall,
+    interpreter: &Interpreter,
     tags: &Tags,
     markers: &MarkerEnvironment,
     client: &RegistryClient,
@@ -361,8 +363,16 @@ async fn resolve(
     );
 
     // Resolve the dependencies.
-    let resolver = Resolver::new(manifest, options, markers, tags, client, build_dispatch)
-        .with_reporter(ResolverReporter::from(printer));
+    let resolver = Resolver::new(
+        manifest,
+        options,
+        markers,
+        interpreter,
+        tags,
+        client,
+        build_dispatch,
+    )
+    .with_reporter(ResolverReporter::from(printer));
     let resolution = resolver.resolve().await?;
 
     let s = if resolution.len() == 1 { "" } else { "s" };

crates/puffin-cli/tests/pip_compile.rs

@@ -669,8 +669,9 @@ fn compile_python_37() -> Result<()> {
 
         ----- stderr -----
           × No solution found when resolving dependencies:
-          ╰─▶ Because there is no version of black available matching ==23.10.1 and
-              root depends on black==23.10.1, version solving failed.
+          ╰─▶ Because there is no version of Python available matching >=3.8 and
+              black==23.10.1 depends on Python>=3.8, black==23.10.1 is forbidden.
+              And because root depends on black==23.10.1, version solving failed.
         "###);
     });
 
@@ -1552,8 +1553,8 @@ fn conflicting_transitive_url_dependency() -> Result<()> {
 
         ----- stderr -----
           × No solution found when resolving dependencies:
-          ╰─▶ Because there is no version of werkzeug available matching >=3.0.0 and
-              flask==3.0.0 depends on werkzeug>=3.0.0, flask==3.0.0 is forbidden.
+          ╰─▶ Because flask==3.0.0 depends on werkzeug>=3.0.0 and there is no version
+              of werkzeug available matching >=3.0.0, flask==3.0.0 is forbidden.
               And because root depends on flask==3.0.0, version solving failed.
         "###);
     });

crates/puffin-cli/tests/pip_install.rs

@@ -74,11 +74,11 @@ fn no_solution() -> Result<()> {
 
     ----- stderr -----
       × No solution found when resolving dependencies:
-      ╰─▶ Because flask==3.0.0 depends on werkzeug>=3.0.0 and there is no
-          version of flask available matching >3.0.0, flask>=3.0.0 depends on
+      ╰─▶ Because there is no version of flask available matching >3.0.0 and
+          flask==3.0.0 depends on werkzeug>=3.0.0, flask>=3.0.0 depends on
           werkzeug>=3.0.0.
-          And because root depends on werkzeug<1.0.0 and root depends on
-          flask>=3.0.0, version solving failed.
+          And because root depends on flask>=3.0.0 and root depends on
+          werkzeug<1.0.0, version solving failed.
     "###);
 
     Ok(())

crates/puffin-dev/src/resolve_cli.rs

@@ -66,6 +66,7 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> {
         Manifest::simple(args.requirements.clone()),
         ResolutionOptions::default(),
         venv.interpreter().markers(),
+        venv.interpreter(),
         tags,
         &client,
         &build_dispatch,

crates/puffin-dispatch/src/lib.rs

@@ -81,7 +81,6 @@ impl<'a> BuildContext for BuildDispatch<'a> {
         self.no_build
     }
 
-    //#[instrument(skip(self, requirements), fields(requirements = requirements.iter().map(ToString::to_string).join(", ")))]
     async fn resolve<'data>(&'data self, requirements: &'data [Requirement]) -> Result<Resolution> {
         let markers = self.interpreter.markers();
         let tags = self.interpreter.tags()?;
@@ -89,6 +88,7 @@ impl<'a> BuildContext for BuildDispatch<'a> {
             Manifest::simple(requirements.to_vec()),
             self.options,
             markers,
+            self.interpreter,
             tags,
             self.client,
             self,

crates/puffin-resolver/src/candidate_selector.rs

@@ -2,6 +2,7 @@ use pubgrub::range::Range;
 use rustc_hash::FxHashMap;
 
 use distribution_types::{Dist, DistributionMetadata, IndexUrl, Name};
+use pep440_rs::VersionSpecifiers;
 use pep508_rs::{Requirement, VersionOrUrl};
 use puffin_normalize::PackageName;
 use pypi_types::BaseUrl;
@@ -9,6 +10,7 @@ use pypi_types::BaseUrl;
 use crate::file::DistFile;
 use crate::prerelease_mode::PreReleaseStrategy;
 use crate::pubgrub::PubGrubVersion;
+use crate::python_requirement::PythonRequirement;
 use crate::resolution_mode::ResolutionStrategy;
 use crate::version_map::{ResolvableFile, VersionMap};
 use crate::{Manifest, ResolutionOptions};
@@ -254,6 +256,30 @@ impl<'a> Candidate<'a> {
         self.file.install()
     }
 
+    /// If the candidate doesn't the given requirement, return the version specifiers.
+    pub(crate) fn validate(&self, requirement: &PythonRequirement) -> Option<&VersionSpecifiers> {
+        // Validate against the _installed_ file. It's fine if the _resolved_ file is incompatible,
+        // since it could be an incompatible wheel. (If the resolved file is an incompatible source
+        // distribution, then the resolved and installed file will be the same anyway.)
+        let requires_python = self.install().requires_python.as_ref()?;
+
+        // If the candidate doesn't support the target Python version, return the failing version
+        // specifiers.
+        if !requires_python.contains(requirement.target()) {
+            return Some(requires_python);
+        }
+
+        // If the candidate is a source distribution, and doesn't support the installed Python
+        // version, return the failing version specifiers, since we won't be able to build it.
+        if self.install().is_sdist() {
+            if !requires_python.contains(requirement.installed()) {
+                return Some(requires_python);
+            }
+        }
+
+        None
+    }
+
     /// Return the [`Dist`] to use when resolving the candidate.
     pub(crate) fn into_distribution(self, index: IndexUrl, base: BaseUrl) -> Dist {
         Dist::from_registry(

crates/puffin-resolver/src/error.rs

@@ -15,7 +15,10 @@ use puffin_traits::OnceMap;
 use pypi_types::BaseUrl;
 
 use crate::candidate_selector::CandidateSelector;
-use crate::pubgrub::{PubGrubHints, PubGrubPackage, PubGrubReportFormatter, PubGrubVersion};
+use crate::pubgrub::{
+    PubGrubHints, PubGrubPackage, PubGrubPython, PubGrubReportFormatter, PubGrubVersion,
+};
+use crate::python_requirement::PythonRequirement;
 use crate::version_map::VersionMap;
 
 #[derive(Error, Debug)]
@@ -159,21 +162,37 @@ impl NoSolutionError {
     #[must_use]
     pub(crate) fn with_available_versions(
         mut self,
+        python_requirement: &PythonRequirement,
         package_versions: &OnceMap<PackageName, (IndexUrl, BaseUrl, VersionMap)>,
     ) -> Self {
         let mut available_versions = FxHashMap::default();
         for package in self.derivation_tree.packages() {
-            if let PubGrubPackage::Package(name, ..) = package {
-                if let Some(entry) = package_versions.get(name) {
-                    let (_, _, version_map) = entry.value();
+            match package {
+                PubGrubPackage::Root(_) => {}
+                PubGrubPackage::Python(PubGrubPython::Installed) => {
                     available_versions.insert(
                         package.clone(),
-                        version_map
-                            .iter()
-                            .map(|(version, _)| version.clone())
-                            .collect(),
+                        vec![PubGrubVersion::from(python_requirement.installed().clone())],
                     );
                 }
+                PubGrubPackage::Python(PubGrubPython::Target) => {
+                    available_versions.insert(
+                        package.clone(),
+                        vec![PubGrubVersion::from(python_requirement.target().clone())],
+                    );
+                }
+                PubGrubPackage::Package(name, ..) => {
+                    if let Some(entry) = package_versions.get(name) {
+                        let (_, _, version_map) = entry.value();
+                        available_versions.insert(
+                            package.clone(),
+                            version_map
+                                .iter()
+                                .map(|(version, _)| version.clone())
+                                .collect(),
+                        );
+                    }
+                }
             }
         }
         self.available_versions = available_versions;

crates/puffin-resolver/src/file.rs

@@ -62,6 +62,13 @@ impl DistFile {
             Self::Sdist(sdist) => sdist.filename.as_str(),
         }
     }
+
+    pub(crate) fn is_sdist(&self) -> bool {
+        match self {
+            Self::Wheel(_) => false,
+            Self::Sdist(_) => true,
+        }
+    }
 }
 
 impl From<DistFile> for File {

crates/puffin-resolver/src/finder.rs

@@ -134,7 +134,6 @@ impl<'a> DistFinder<'a> {
                 // This is relevant for source dists which give no other indication of their
                 // compatibility and wheels which may be tagged `py3-none-any` but
                 // have `requires-python: ">=3.9"`
-                // TODO(konstin): https://github.com/astral-sh/puffin/issues/406
                 if !file
                     .requires_python
                     .as_ref()
@@ -170,7 +169,6 @@ impl<'a> DistFinder<'a> {
                     // This is relevant for source dists which give no other indication of their
                     // compatibility and wheels which may be tagged `py3-none-any` but
                     // have `requires-python: ">=3.9"`
-                    // TODO(konstin): https://github.com/astral-sh/puffin/issues/406
                     if !file
                         .requires_python
                         .as_ref()

crates/puffin-resolver/src/lib.rs

@@ -18,6 +18,7 @@ mod overrides;
 mod pins;
 mod prerelease_mode;
 mod pubgrub;
+mod python_requirement;
 mod resolution;
 mod resolution_mode;
 mod resolution_options;

crates/puffin-resolver/src/pubgrub/dependencies.rs

@@ -206,6 +206,9 @@ fn merge_package(
         // Either package is `root`.
         (PubGrubPackage::Root(_), _) | (_, PubGrubPackage::Root(_)) => Ok(None),
 
+        // Either package is the Python installation.
+        (PubGrubPackage::Python(_), _) | (_, PubGrubPackage::Python(_)) => Ok(None),
+
         // Left package has a URL. Propagate the URL.
         (PubGrubPackage::Package(name, extra, Some(url)), PubGrubPackage::Package(.., None)) => {
             Ok(Some(PubGrubPackage::Package(

crates/puffin-resolver/src/pubgrub/mod.rs

@@ -1,8 +1,9 @@
 pub(crate) use crate::pubgrub::dependencies::PubGrubDependencies;
 pub(crate) use crate::pubgrub::distribution::PubGrubDistribution;
-pub(crate) use crate::pubgrub::package::PubGrubPackage;
+pub(crate) use crate::pubgrub::package::{PubGrubPackage, PubGrubPython};
 pub(crate) use crate::pubgrub::priority::{PubGrubPriorities, PubGrubPriority};
 pub(crate) use crate::pubgrub::report::{PubGrubHints, PubGrubReportFormatter};
+pub(crate) use crate::pubgrub::specifier::PubGrubSpecifier;
 pub(crate) use crate::pubgrub::version::{PubGrubVersion, MIN_VERSION};
 
 mod dependencies;

crates/puffin-resolver/src/pubgrub/package.rs

@@ -13,7 +13,11 @@ use puffin_normalize::{ExtraName, PackageName};
 #[derive(Debug, Clone, Eq, Derivative)]
 #[derivative(PartialEq, Hash)]
 pub enum PubGrubPackage {
+    /// The root package, which is used to start the resolution process.
     Root(Option<PackageName>),
+    /// A Python version.
+    Python(PubGrubPython),
+    /// A Python package.
     Package(
         PackageName,
         Option<ExtraName>,
@@ -70,6 +74,14 @@ pub enum PubGrubPackage {
     ),
 }
 
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+pub enum PubGrubPython {
+    /// The Python version installed in the current environment.
+    Installed,
+    /// The Python version for which dependencies are being resolved.
+    Target,
+}
+
 impl std::fmt::Display for PubGrubPackage {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -80,6 +92,7 @@ impl std::fmt::Display for PubGrubPackage {
                     write!(f, "root")
                 }
             }
+            PubGrubPackage::Python(_) => write!(f, "Python"),
             PubGrubPackage::Package(name, None, ..) => write!(f, "{name}"),
             PubGrubPackage::Package(name, Some(extra), ..) => {
                 write!(f, "{name}[{extra}]")

crates/puffin-resolver/src/pubgrub/priority.rs

@@ -20,6 +20,7 @@ impl PubGrubPriorities {
     pub(crate) fn get(&self, package: &PubGrubPackage) -> Option<PubGrubPriority> {
         match package {
             PubGrubPackage::Root(_) => Some(Reverse(0)),
+            PubGrubPackage::Python(_) => Some(Reverse(0)),
             PubGrubPackage::Package(name, _, _) => self
                 .0
                 .get(name)

crates/puffin-resolver/src/python_requirement.rs

@@ -0,0 +1,37 @@
+use pep440_rs::Version;
+use pep508_rs::MarkerEnvironment;
+use puffin_interpreter::Interpreter;
+
+#[derive(Debug, Clone)]
+pub struct PythonRequirement<'a> {
+    /// The installed version of Python.
+    installed: &'a Version,
+    /// The target version of Python; that is, the version of Python for which we are resolving
+    /// dependencies. This is typically the same as the installed version, but may be different
+    /// when specifying an alternate Python version for the resolution.
+    target: &'a Version,
+}
+
+impl<'a> PythonRequirement<'a> {
+    pub fn new(interpreter: &'a Interpreter, markers: &'a MarkerEnvironment) -> Self {
+        Self {
+            installed: interpreter.version(),
+            target: &markers.python_version.version,
+        }
+    }
+
+    /// Return the installed version of Python.
+    pub(crate) fn installed(&self) -> &'a Version {
+        self.installed
+    }
+
+    /// Return the target version of Python.
+    pub(crate) fn target(&self) -> &'a Version {
+        self.target
+    }
+
+    /// Returns an iterator over the versions of Python to consider when resolving dependencies.
+    pub(crate) fn versions(&self) -> impl Iterator<Item = &'a Version> {
+        std::iter::once(self.installed).chain(std::iter::once(self.target))
+    }
+}

crates/puffin-resolver/src/resolution_options.rs

@@ -1,10 +1,10 @@
-use crate::{PreReleaseMode, ResolutionMode};
 use chrono::{DateTime, Utc};
 
+use crate::{PreReleaseMode, ResolutionMode};
+
 /// Options for resolving a manifest.
 #[derive(Debug, Default, Copy, Clone)]
 pub struct ResolutionOptions {
-    // TODO(konstin): These should be pub(crate) again
     pub resolution_mode: ResolutionMode,
     pub prerelease_mode: PreReleaseMode,
     pub exclude_newer: Option<DateTime<Utc>>,

crates/puffin-resolver/src/resolver.rs

@@ -7,6 +7,7 @@ use anyhow::Result;
 use chrono::{DateTime, Utc};
 use futures::channel::mpsc::UnboundedReceiver;
 use futures::{pin_mut, FutureExt, StreamExt, TryFutureExt};
+use itertools::Itertools;
 use pubgrub::error::PubGrubError;
 use pubgrub::range::Range;
 use pubgrub::solver::{Incompatibility, State};
@@ -21,10 +22,12 @@ use distribution_types::{
     BuiltDist, Dist, DistributionMetadata, IndexUrl, LocalEditable, Name, PackageId, SourceDist,
     VersionOrUrl,
 };
+use pep440_rs::VersionSpecifiers;
 use pep508_rs::{MarkerEnvironment, Requirement};
 use platform_tags::Tags;
 use puffin_client::RegistryClient;
 use puffin_distribution::{DistributionDatabase, DistributionDatabaseError};
+use puffin_interpreter::Interpreter;
 use puffin_normalize::PackageName;
 use puffin_traits::{BuildContext, OnceMap};
 use pypi_types::{BaseUrl, Metadata21};
@@ -35,9 +38,10 @@ use crate::manifest::Manifest;
 use crate::overrides::Overrides;
 use crate::pins::FilePins;
 use crate::pubgrub::{
-    PubGrubDependencies, PubGrubDistribution, PubGrubPackage, PubGrubPriorities, PubGrubVersion,
-    MIN_VERSION,
+    PubGrubDependencies, PubGrubDistribution, PubGrubPackage, PubGrubPriorities, PubGrubPython,
+    PubGrubSpecifier, PubGrubVersion, MIN_VERSION,
 };
+use crate::python_requirement::PythonRequirement;
 use crate::resolution::ResolutionGraph;
 use crate::version_map::VersionMap;
 use crate::yanks::AllowedYanks;
@@ -73,9 +77,8 @@ pub trait ResolverProvider: Send + Sync {
 pub struct DefaultResolverProvider<'a, Context: BuildContext + Send + Sync> {
     client: &'a RegistryClient,
     fetcher: DistributionDatabase<'a, Context>,
-    build_context: &'a Context,
     tags: &'a Tags,
-    markers: &'a MarkerEnvironment,
+    python_requirement: PythonRequirement<'a>,
     exclude_newer: Option<DateTime<Utc>>,
     allowed_yanks: AllowedYanks,
 }
@@ -84,18 +87,16 @@ impl<'a, Context: BuildContext + Send + Sync> DefaultResolverProvider<'a, Contex
     pub fn new(
         client: &'a RegistryClient,
         fetcher: DistributionDatabase<'a, Context>,
-        build_context: &'a Context,
         tags: &'a Tags,
-        markers: &'a MarkerEnvironment,
+        python_requirement: PythonRequirement<'a>,
         exclude_newer: Option<DateTime<Utc>>,
         allowed_yanks: AllowedYanks,
     ) -> Self {
         Self {
             client,
             fetcher,
-            build_context,
             tags,
-            markers,
+            python_requirement,
             exclude_newer,
             allowed_yanks,
         }
@@ -119,8 +120,7 @@ impl<'a, Context: BuildContext + Send + Sync> ResolverProvider
                         metadata,
                         package_name,
                         self.tags,
-                        self.markers,
-                        self.build_context.interpreter(),
+                        &self.python_requirement,
                         &self.allowed_yanks,
                         self.exclude_newer.as_ref(),
                     ),
@@ -152,6 +152,7 @@ pub struct Resolver<'a, Provider: ResolverProvider> {
     overrides: Overrides,
     allowed_urls: AllowedUrls,
     markers: &'a MarkerEnvironment,
+    python_requirement: PythonRequirement<'a>,
     selector: CandidateSelector,
     index: Arc<Index>,
     editables: FxHashMap<PackageName, (LocalEditable, Metadata21)>,
@@ -165,6 +166,7 @@ impl<'a, Context: BuildContext + Send + Sync> Resolver<'a, DefaultResolverProvid
         manifest: Manifest,
         options: ResolutionOptions,
         markers: &'a MarkerEnvironment,
+        interpreter: &'a Interpreter,
         tags: &'a Tags,
         client: &'a RegistryClient,
         build_context: &'a Context,
@@ -172,9 +174,8 @@ impl<'a, Context: BuildContext + Send + Sync> Resolver<'a, DefaultResolverProvid
         let provider = DefaultResolverProvider::new(
             client,
             DistributionDatabase::new(build_context.cache(), tags, client, build_context),
-            build_context,
             tags,
-            markers,
+            PythonRequirement::new(interpreter, markers),
             options.exclude_newer,
             manifest
                 .requirements
@@ -182,7 +183,13 @@ impl<'a, Context: BuildContext + Send + Sync> Resolver<'a, DefaultResolverProvid
                 .chain(manifest.constraints.iter())
                 .collect(),
         );
-        Self::new_custom_io(manifest, options, markers, provider)
+        Self::new_custom_io(
+            manifest,
+            options,
+            markers,
+            PythonRequirement::new(interpreter, markers),
+            provider,
+        )
     }
 }
 
@@ -192,6 +199,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
         manifest: Manifest,
         options: ResolutionOptions,
         markers: &'a MarkerEnvironment,
+        python_requirement: PythonRequirement<'a>,
         provider: Provider,
     ) -> Self {
         let selector = CandidateSelector::for_resolution(&manifest, options);
@@ -245,6 +253,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
             constraints: manifest.constraints,
             overrides: Overrides::from_requirements(manifest.overrides),
             markers,
+            python_requirement,
             editables,
             reporter: None,
             provider,
@@ -287,7 +296,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                 resolution.map_err(|err| {
                     // Add version information to improve unsat error messages
                     if let ResolveError::NoSolution(err) = err {
-                        ResolveError::NoSolution(err.with_available_versions(&self.index.packages).with_selector(self.selector.clone()))
+                        ResolveError::NoSolution(err.with_available_versions(&self.python_requirement, &self.index.packages).with_selector(self.selector.clone()))
                     } else {
                         err
                     }
@@ -440,6 +449,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
     ) -> Result<(), ResolveError> {
         match package {
             PubGrubPackage::Root(_) => {}
+            PubGrubPackage::Python(_) => {}
             PubGrubPackage::Package(package_name, _extra, None) => {
                 // Emit a request to fetch the metadata for this package.
                 if index.packages.register(package_name) {
@@ -488,6 +498,24 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
         return match package {
             PubGrubPackage::Root(_) => Ok(Some(MIN_VERSION.clone())),
 
+            PubGrubPackage::Python(PubGrubPython::Installed) => {
+                let version = PubGrubVersion::from(self.python_requirement.installed().clone());
+                if range.contains(&version) {
+                    Ok(Some(version))
+                } else {
+                    Ok(None)
+                }
+            }
+
+            PubGrubPackage::Python(PubGrubPython::Target) => {
+                let version = PubGrubVersion::from(self.python_requirement.target().clone());
+                if range.contains(&version) {
+                    Ok(Some(version))
+                } else {
+                    Ok(None)
+                }
+            }
+
             PubGrubPackage::Package(package_name, extra, Some(url)) => {
                 if let Some(extra) = extra {
                     debug!(
@@ -548,6 +576,14 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                     return Ok(None);
                 };
 
+                // If the version is incompatible, short-circuit.
+                if let Some(requires_python) = candidate.validate(&self.python_requirement) {
+                    self.index
+                        .incompatibilities
+                        .done(candidate.package_id(), requires_python.clone());
+                    return Ok(Some(candidate.version().clone()));
+                }
+
                 if let Some(extra) = extra {
                     debug!(
                         "Selecting: {}[{}]=={} ({})",
@@ -636,13 +672,37 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                 Ok(Dependencies::Known(constraints.into()))
             }
 
+            PubGrubPackage::Python(_) => Ok(Dependencies::Known(DependencyConstraints::default())),
+
             PubGrubPackage::Package(package_name, extra, url) => {
                 // Wait for the metadata to be available.
                 let dist = match url {
                     Some(url) => PubGrubDistribution::from_url(package_name, url),
                     None => PubGrubDistribution::from_registry(package_name, version),
                 };
-                let entry = self.index.distributions.wait(&dist.package_id()).await;
+                let package_id = dist.package_id();
+
+                // If the package is known to be incompatible, return the Python version as an
+                // incompatibility, and skip fetching the metadata.
+                if let Some(entry) = self.index.incompatibilities.get(&package_id) {
+                    let requires_python = entry.value();
+                    let version = requires_python
+                        .iter()
+                        .map(PubGrubSpecifier::try_from)
+                        .fold_ok(Range::full(), |range, specifier| {
+                            range.intersection(&specifier.into())
+                        })?;
+
+                    let mut constraints = DependencyConstraints::default();
+                    constraints.insert(
+                        PubGrubPackage::Python(PubGrubPython::Installed),
+                        version.clone(),
+                    );
+                    constraints.insert(PubGrubPackage::Python(PubGrubPython::Target), version);
+                    return Ok(Dependencies::Known(constraints));
+                }
+
+                let entry = self.index.distributions.wait(&package_id).await;
                 let metadata = entry.value();
 
                 let mut constraints = PubGrubDependencies::from_requirements(
@@ -661,6 +721,23 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                     Self::visit_package(package, priorities, index, request_sink)?;
                 }
 
+                // If a package has a `requires_python` field, add a constraint on the target
+                // Python version.
+                if let Some(requires_python) = metadata.requires_python.as_ref() {
+                    let version = requires_python
+                        .iter()
+                        .map(PubGrubSpecifier::try_from)
+                        .fold_ok(Range::full(), |range, specifier| {
+                            range.intersection(&specifier.into())
+                        })?;
+                    constraints.insert(
+                        PubGrubPackage::Python(PubGrubPython::Installed),
+                        version.clone(),
+                    );
+                    constraints.insert(PubGrubPackage::Python(PubGrubPython::Target), version);
+                }
+
+                // If a package has an extra, insert a constraint on the base package.
                 if extra.is_some() {
                     constraints.insert(
                         PubGrubPackage::Package(package_name.clone(), None, None),
@@ -766,6 +843,14 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                     return Ok(None);
                 };
 
+                // If the version is incompatible, short-circuit.
+                if let Some(requires_python) = candidate.validate(&self.python_requirement) {
+                    self.index
+                        .incompatibilities
+                        .done(candidate.package_id(), requires_python.clone());
+                    return Ok(None);
+                }
+
                 // Emit a request to fetch the metadata for this version.
                 if self.index.distributions.register(&candidate.package_id()) {
                     let dist = candidate.into_distribution(index.clone(), base.clone());
@@ -802,6 +887,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
         if let Some(reporter) = self.reporter.as_ref() {
             match package {
                 PubGrubPackage::Root(_) => {}
+                PubGrubPackage::Python(_) => {}
                 PubGrubPackage::Package(package_name, _extra, Some(url)) => {
                     reporter.on_progress(package_name, VersionOrUrl::Url(url));
                 }
@@ -892,9 +978,12 @@ pub(crate) struct Index {
     /// came from.
     pub(crate) packages: OnceMap<PackageName, (IndexUrl, BaseUrl, VersionMap)>,
 
-    /// A map from distribution SHA to metadata for that distribution.
+    /// A map from package ID to metadata for that distribution.
     pub(crate) distributions: OnceMap<PackageId, Metadata21>,
 
+    /// A map from package ID to required Python version.
+    pub(crate) incompatibilities: OnceMap<PackageId, VersionSpecifiers>,
+
     /// A map from source URL to precise URL.
     pub(crate) redirects: OnceMap<Url, Url>,
 }

crates/puffin-resolver/src/version_map.rs

@@ -5,16 +5,15 @@ use chrono::{DateTime, Utc};
 use tracing::{instrument, warn};
 
 use distribution_filename::DistFilename;
-use pep508_rs::MarkerEnvironment;
 use platform_tags::{TagPriority, Tags};
 use puffin_client::SimpleMetadata;
-use puffin_interpreter::Interpreter;
 use puffin_normalize::PackageName;
 use puffin_warnings::warn_user_once;
 use pypi_types::Yanked;
 
 use crate::file::{DistFile, SdistFile, WheelFile};
 use crate::pubgrub::PubGrubVersion;
+use crate::python_requirement::PythonRequirement;
 use crate::yanks::AllowedYanks;
 
 /// A map from versions to distributions.
@@ -23,13 +22,12 @@ pub struct VersionMap(BTreeMap<PubGrubVersion, PrioritizedDistribution>);
 
 impl VersionMap {
     /// Initialize a [`VersionMap`] from the given metadata.
-    #[instrument(skip_all, fields(package_name = %package_name))]
+    #[instrument(skip_all, fields(package_name = % package_name))]
     pub(crate) fn from_metadata(
         metadata: SimpleMetadata,
         package_name: &PackageName,
         tags: &Tags,
-        markers: &MarkerEnvironment,
-        interpreter: &Interpreter,
+        python_requirement: &PythonRequirement,
         allowed_yanks: &AllowedYanks,
         exclude_newer: Option<&DateTime<Utc>>,
     ) -> Self {
@@ -39,23 +37,6 @@ impl VersionMap {
         // Collect compatible distributions.
         for (version, files) in metadata {
             for (filename, file) in files.all() {
-                // Only add dists compatible with the python version. This is relevant for source
-                // distributions which give no other indication of their compatibility and wheels which
-                // may be tagged `py3-none-any` but have `requires-python: ">=3.9"`.
-                // TODO(konstin): https://github.com/astral-sh/puffin/issues/406
-                if let Some(requires_python) = file.requires_python.as_ref() {
-                    // The interpreter and marker version are often the same, but can differ. For
-                    // example, if the user is resolving against a target Python version passed in
-                    // via the command-line, that version will differ from the interpreter version.
-                    let interpreter_version = interpreter.version();
-                    let marker_version = &markers.python_version.version;
-                    if !requires_python.contains(interpreter_version)
-                        || !requires_python.contains(marker_version)
-                    {
-                        continue;
-                    }
-                }
-
                 // Support resolving as if it were an earlier timestamp, at least as long files have
                 // upload time information
                 if let Some(exclude_newer) = exclude_newer {
@@ -86,8 +67,17 @@ impl VersionMap {
 
                 match filename {
                     DistFilename::WheelFilename(filename) => {
-                        let priority = filename.compatibility(tags);
-
+                        // To be compatible, the wheel must both have compatible tags _and_ have a
+                        // compatible Python requirement.
+                        let priority = filename.compatibility(tags).filter(|_| {
+                            file.requires_python
+                                .as_ref()
+                                .map_or(true, |requires_python| {
+                                    python_requirement
+                                        .versions()
+                                        .all(|version| requires_python.contains(version))
+                                })
+                        });
                         match version_map.entry(version.clone().into()) {
                             Entry::Occupied(mut entry) => {
                                 entry.get_mut().insert_built(WheelFile(file), priority);
@@ -201,12 +191,12 @@ impl PrioritizedDistribution {
         ) {
             // Prefer the highest-priority, platform-compatible wheel.
             (Some((wheel, _)), _, _) => Some(ResolvableFile::CompatibleWheel(wheel)),
-            // If we have a source distribution and an incompatible wheel, return the wheel.
-            // We assume that all distributions have the same metadata for a given package version.
-            // If a source distribution exists, we assume we can build it, but using the wheel is
-            // faster.
+            // If we have a compatible source distribution and an incompatible wheel, return the
+            // wheel. We assume that all distributions have the same metadata for a given package
+            // version. If a compatible source distribution exists, we assume we can build it, but
+            // using the wheel is faster.
             (_, Some(sdist), Some(wheel)) => Some(ResolvableFile::IncompatibleWheel(sdist, wheel)),
-            // Otherwise, return the source distribution.
+            // Otherwise, if we have a source distribution, return it.
             (_, Some(sdist), _) => Some(ResolvableFile::SourceDist(sdist)),
             _ => None,
         }

crates/puffin-resolver/tests/resolver.rs

@@ -91,17 +91,26 @@ async fn resolve(
     tags: &Tags,
 ) -> Result<ResolutionGraph> {
     let client = RegistryClientBuilder::new(Cache::temp()?).build();
+    let interpreter = Interpreter::artificial(
+        Platform::current()?,
+        markers.clone(),
+        PathBuf::from("/dev/null"),
+        PathBuf::from("/dev/null"),
+        PathBuf::from("/dev/null"),
+    );
     let build_context = DummyContext {
         cache: Cache::temp()?,
-        interpreter: Interpreter::artificial(
-            Platform::current()?,
-            markers.clone(),
-            PathBuf::from("/dev/null"),
-            PathBuf::from("/dev/null"),
-            PathBuf::from("/dev/null"),
-        ),
+        interpreter: interpreter.clone(),
     };
-    let resolver = Resolver::new(manifest, options, markers, tags, &client, &build_context);
+    let resolver = Resolver::new(
+        manifest,
+        options,
+        markers,
+        &interpreter,
+        tags,
+        &client,
+        &build_context,
+    );
     Ok(resolver.resolve().await?)
 }