Add support for emitting index URLs and --find-links (#1142)

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

crates/distribution-types/src/index_url.rs

@@ -16,6 +16,15 @@ pub enum IndexUrl {
     Url(Url),
 }
 
+impl Display for IndexUrl {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            IndexUrl::Pypi => Display::fmt(&*PYPI_URL, f),
+            IndexUrl::Url(url) => Display::fmt(url, f),
+        }
+    }
+}
+
 impl FromStr for IndexUrl {
     type Err = url::ParseError;
 
@@ -153,13 +162,23 @@ impl IndexLocations {
 }
 
 impl<'a> IndexLocations {
-    /// Return an iterator over the [`IndexUrl`] entries.
+    /// Return an iterator over all [`IndexUrl`] entries.
     pub fn indexes(&'a self) -> impl Iterator<Item = &'a IndexUrl> + 'a {
         self.index.iter().chain(self.extra_index.iter())
     }
 
+    /// Return the primary [`IndexUrl`] entry.
+    pub fn index(&'a self) -> Option<&'a IndexUrl> {
+        self.index.as_ref()
+    }
+
+    /// Return an iterator over the extra [`IndexUrl`] entries.
+    pub fn extra_index(&'a self) -> impl Iterator<Item = &'a IndexUrl> + 'a {
+        self.extra_index.iter()
+    }
+
     /// Return an iterator over the [`FlatIndexLocation`] entries.
-    pub fn flat_indexes(&'a self) -> impl Iterator<Item = &'a FlatIndexLocation> + 'a {
+    pub fn flat_index(&'a self) -> impl Iterator<Item = &'a FlatIndexLocation> + 'a {
         self.flat_index.iter()
     }
 

crates/puffin-dev/src/resolve_cli.rs

@@ -63,7 +63,7 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> {
         .build();
     let flat_index = {
         let client = FlatIndexClient::new(&client, &cache);
-        let entries = client.fetch(index_locations.flat_indexes()).await?;
+        let entries = client.fetch(index_locations.flat_index()).await?;
         FlatIndex::from_entries(entries, venv.interpreter().tags()?)
     };
     let index = InMemoryIndex::default();

crates/puffin-distribution/src/index/registry_wheel_index.rs

@@ -80,7 +80,7 @@ impl<'a> RegistryWheelIndex<'a> {
 
         // Collect into owned `IndexUrl`
         let flat_index_urls: Vec<IndexUrl> = index_locations
-            .flat_indexes()
+            .flat_index()
             .filter_map(|flat_index| match flat_index {
                 FlatIndexLocation::Path(_) => None,
                 FlatIndexLocation::Url(url) => Some(IndexUrl::Url(url.clone())),

crates/puffin/src/commands/pip_compile.rs

@@ -55,6 +55,8 @@ pub(crate) async fn pip_compile(
     generate_hashes: bool,
     include_annotations: bool,
     include_header: bool,
+    include_index_url: bool,
+    include_find_links: bool,
     index_locations: IndexLocations,
     setup_py: SetupPyStrategy,
     no_build: bool,
@@ -190,7 +192,7 @@ pub(crate) async fn pip_compile(
     // Resolve the flat indexes from `--find-links`.
     let flat_index = {
         let client = FlatIndexClient::new(&client, &cache);
-        let entries = client.fetch(index_locations.flat_indexes()).await?;
+        let entries = client.fetch(index_locations.flat_index()).await?;
         FlatIndex::from_entries(entries, &tags)
     };
 
@@ -342,6 +344,34 @@ pub(crate) async fn pip_compile(
         )?;
     }
 
+    // Write the index locations to the output channel.
+    let mut wrote_index = false;
+
+    // If necessary, include the `--index-url` and `--extra-index-url` locations.
+    if include_index_url {
+        if let Some(index) = index_locations.index() {
+            writeln!(writer, "--index-url {index}")?;
+            wrote_index = true;
+        }
+        for extra_index in index_locations.extra_index() {
+            writeln!(writer, "--extra-index-url {extra_index}")?;
+            wrote_index = true;
+        }
+    }
+
+    // If necessary, include the `--find-links` locations.
+    if include_find_links {
+        for flat_index in index_locations.flat_index() {
+            writeln!(writer, "--find-links {flat_index}")?;
+            wrote_index = true;
+        }
+    }
+
+    // If we wrote an index, add a newline to separate it from the requirements
+    if wrote_index {
+        writeln!(writer)?;
+    }
+
     write!(
         writer,
         "{}",

crates/puffin/src/commands/pip_install.rs

@@ -133,7 +133,7 @@ pub(crate) async fn pip_install(
     // Resolve the flat indexes from `--find-links`.
     let flat_index = {
         let client = FlatIndexClient::new(&client, &cache);
-        let entries = client.fetch(index_locations.flat_indexes()).await?;
+        let entries = client.fetch(index_locations.flat_index()).await?;
         FlatIndex::from_entries(entries, tags)
     };
 

crates/puffin/src/commands/pip_sync.rs

@@ -73,7 +73,7 @@ pub(crate) async fn pip_sync(
     // Resolve the flat indexes from `--find-links`.
     let flat_index = {
         let client = FlatIndexClient::new(&client, &cache);
-        let entries = client.fetch(index_locations.flat_indexes()).await?;
+        let entries = client.fetch(index_locations.flat_index()).await?;
         FlatIndex::from_entries(entries, tags)
     };
 
@@ -167,7 +167,7 @@ pub(crate) async fn pip_sync(
         // Resolve the flat indexes from `--find-links`.
         let flat_index = {
             let client = FlatIndexClient::new(&client, &cache);
-            let entries = client.fetch(index_locations.flat_indexes()).await?;
+            let entries = client.fetch(index_locations.flat_index()).await?;
             FlatIndex::from_entries(entries, tags)
         };
 

crates/puffin/src/commands/venv.rs

@@ -118,7 +118,7 @@ async fn venv_impl(
             let tags = interpreter.tags().map_err(VenvError::Tags)?;
             let client = FlatIndexClient::new(&client, cache);
             let entries = client
-                .fetch(index_locations.flat_indexes())
+                .fetch(index_locations.flat_index())
                 .await
                 .map_err(VenvError::FlatIndex)?;
             FlatIndex::from_entries(entries, tags)

crates/puffin/src/compat/mod.rs

@@ -59,15 +59,9 @@ pub(crate) struct PipCompileCompatArgs {
     #[clap(long, hide = true)]
     no_config: bool,
 
-    #[clap(long, hide = true)]
-    emit_index_url: bool,
-
     #[clap(long, hide = true)]
     no_emit_index_url: bool,
 
-    #[clap(long, hide = true)]
-    emit_find_links: bool,
-
     #[clap(long, hide = true)]
     no_emit_find_links: bool,
 
@@ -195,27 +189,15 @@ impl PipCompileCompatArgs {
             );
         }
 
-        if self.emit_index_url {
-            return Err(anyhow!(
-                "pip-compile's `--emit-index-url` is unsupported (Puffin never emits index URLs)."
-            ));
-        }
-
         if self.no_emit_index_url {
             warn_user!(
-                "pip-compile's `--no-emit-index-url` has no effect (Puffin never emits index URLs)."
+                "pip-compile's `--no-emit-index-url` has no effect (Puffin excludes index URLs by default)."
             );
         }
 
-        if self.emit_find_links {
-            return Err(anyhow!(
-                "pip-compile's `--emit-find-links` is unsupported (Puffin never emits `--find-links` URLs)."
-            ));
-        }
-
         if self.no_emit_find_links {
             warn_user!(
-                "pip-compile's `--no-emit-find-links` has no effect (Puffin never emits `--find-links` URLs)."
+                "pip-compile's `--no-emit-find-links` has no effect (Puffin excludes `--find-links` URLs by default)."
             );
         }
 

crates/puffin/src/main.rs

@@ -225,6 +225,10 @@ struct PipCompileArgs {
     #[clap(long)]
     extra_index_url: Vec<IndexUrl>,
 
+    /// Ignore the package index, instead relying on local archives and caches.
+    #[clap(long, conflicts_with = "index_url", conflicts_with = "extra_index_url")]
+    no_index: bool,
+
     /// Locations to search for candidate distributions, beyond those found in the indexes.
     ///
     /// If a path, the target must be a directory that contains package as wheel files (`.whl`) or
@@ -234,10 +238,6 @@ struct PipCompileArgs {
     #[clap(long)]
     find_links: Vec<FlatIndexLocation>,
 
-    /// Ignore the package index, instead relying on local archives and caches.
-    #[clap(long, conflicts_with = "index_url", conflicts_with = "extra_index_url")]
-    no_index: bool,
-
     /// Allow package upgrades, ignoring pinned versions in the existing output file.
     #[clap(long)]
     upgrade: bool,
@@ -284,6 +284,14 @@ struct PipCompileArgs {
     #[arg(long, value_parser = date_or_datetime)]
     exclude_newer: Option<DateTime<Utc>>,
 
+    /// Include `--index-url` and `--extra-index-url` entries in the generated output file.
+    #[clap(long, hide = true)]
+    emit_index_url: bool,
+
+    /// Include `--find-links` entries in the generated output file.
+    #[clap(long, hide = true)]
+    emit_find_links: bool,
+
     #[command(flatten)]
     compat_args: compat::PipCompileCompatArgs,
 }
@@ -695,6 +703,8 @@ async fn inner() -> Result<ExitStatus> {
                 args.generate_hashes,
                 !args.no_annotate,
                 !args.no_header,
+                args.emit_index_url,
+                args.emit_find_links,
                 index_urls,
                 if args.legacy_setup_py {
                     SetupPyStrategy::Setuptools

crates/puffin/tests/pip_compile.rs

@@ -3867,3 +3867,110 @@ fn resolver_legacy() -> Result<()> {
 
     Ok(())
 }
+
+/// Emit the `--index-url` and `--extra-index-url` locations.
+#[test]
+fn emit_index_urls() -> Result<()> {
+    let temp_dir = TempDir::new()?;
+    let cache_dir = TempDir::new()?;
+    let venv = create_venv_py312(&temp_dir, &cache_dir);
+
+    let requirements_in = temp_dir.child("requirements.in");
+    requirements_in.write_str("black==23.10.1")?;
+
+    insta::with_settings!({
+        filters => INSTA_FILTERS.to_vec()
+    }, {
+        assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME))
+            .arg("pip")
+            .arg("compile")
+            .arg("requirements.in")
+            .arg("--emit-index-url")
+            .arg("--extra-index-url")
+            .arg("https://test.pypi.org/simple/")
+            .arg("--cache-dir")
+            .arg(cache_dir.path())
+            .arg("--exclude-newer")
+            .arg(EXCLUDE_NEWER)
+            .env("VIRTUAL_ENV", venv.as_os_str())
+            .current_dir(&temp_dir), @r###"
+        success: true
+        exit_code: 0
+        ----- stdout -----
+        # This file was autogenerated by Puffin v[VERSION] via the following command:
+        #    puffin pip compile requirements.in --emit-index-url --extra-index-url https://test.pypi.org/simple/ --cache-dir [CACHE_DIR]
+        --index-url https://pypi.org/simple
+        --extra-index-url https://test.pypi.org/simple/
+
+        black==23.10.1
+        click==8.1.7
+            # via black
+        mypy-extensions==1.0.0
+            # via black
+        packaging==23.2
+            # via black
+        pathspec==0.11.2
+            # via black
+        platformdirs==4.0.0
+            # via black
+
+        ----- stderr -----
+        Resolved 6 packages in [TIME]
+        "###);
+    });
+
+    Ok(())
+}
+
+/// Emit the `--find-links` locations.
+#[test]
+fn emit_find_links() -> Result<()> {
+    let temp_dir = TempDir::new()?;
+    let cache_dir = TempDir::new()?;
+    let venv = create_venv_py312(&temp_dir, &cache_dir);
+
+    let requirements_in = temp_dir.child("requirements.in");
+    requirements_in.write_str("black==23.10.1")?;
+
+    insta::with_settings!({
+        filters => INSTA_FILTERS.to_vec()
+    }, {
+        assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME))
+            .arg("pip")
+            .arg("compile")
+            .arg("requirements.in")
+            .arg("--emit-find-links")
+            .arg("--find-links")
+            .arg("./")
+            .arg("--cache-dir")
+            .arg(cache_dir.path())
+            .arg("--exclude-newer")
+            .arg(EXCLUDE_NEWER)
+            .env("VIRTUAL_ENV", venv.as_os_str())
+            .current_dir(&temp_dir), @r###"
+        success: true
+        exit_code: 0
+        ----- stdout -----
+        # This file was autogenerated by Puffin v[VERSION] via the following command:
+        #    puffin pip compile requirements.in --emit-find-links --find-links ./ --cache-dir [CACHE_DIR]
+        --find-links ./
+
+        black==23.10.1
+        click==8.1.7
+            # via black
+        mypy-extensions==1.0.0
+            # via black
+        packaging==23.2
+            # via black
+        pathspec==0.11.2
+            # via black
+        platformdirs==4.0.0
+            # via black
+
+        ----- stderr -----
+        Resolved 6 packages in [TIME]
+        "###);
+    });
+
+    Ok(())
+}