Support extras in editable requirements (#1113)

## Summary

This PR adds support for requirements like `-e .[d]`.

Closes #1091.

crates/distribution-types/src/editable.rs

@@ -4,6 +4,7 @@ use std::path::PathBuf;
 use url::Url;
 
 use pep508_rs::VerbatimUrl;
+use puffin_normalize::ExtraName;
 
 use crate::Verbatim;
 
@@ -13,6 +14,8 @@ pub struct LocalEditable {
     pub url: VerbatimUrl,
     /// Either the path to the editable or its checkout.
     pub path: PathBuf,
+    /// The extras that should be installed.
+    pub extras: Vec<ExtraName>,
 }
 
 impl LocalEditable {

crates/distribution-types/src/lib.rs

@@ -278,7 +278,7 @@ impl Dist {
 
     /// Create a [`Dist`] for a local editable distribution.
     pub fn from_editable(name: PackageName, editable: LocalEditable) -> Result<Self, Error> {
-        let LocalEditable { url, path } = editable;
+        let LocalEditable { url, path, .. } = editable;
         Ok(Self::Source(SourceDist::Path(PathSourceDist {
             name,
             url,

crates/pep508-rs/src/lib.rs

@@ -407,6 +407,22 @@ impl Requirement {
     }
 }
 
+/// A list of [`ExtraName`] that can be attached to a [`Requirement`].
+#[derive(Debug, Clone, Eq, Hash, PartialEq)]
+pub struct Extras(Vec<ExtraName>);
+
+impl Extras {
+    /// Parse a list of extras.
+    pub fn parse(input: &str) -> Result<Self, Pep508Error> {
+        Ok(Self(parse_extras(&mut Cursor::new(input))?))
+    }
+
+    /// Convert the [`Extras`] into a [`Vec`] of [`ExtraName`].
+    pub fn into_vec(self) -> Vec<ExtraName> {
+        self.0
+    }
+}
+
 /// The actual version specifier or url to install
 #[derive(Debug, Clone, Eq, Hash, PartialEq)]
 pub enum VersionOrUrl {

crates/puffin-resolver/src/resolution.rs

@@ -162,14 +162,19 @@ impl ResolutionGraph {
                     dependency_range,
                 ) = &state.incompatibility_store[*id].kind
                 {
-                    let PubGrubPackage::Package(self_package, None, _) = self_package else {
+                    let PubGrubPackage::Package(self_package, _, _) = self_package else {
                         continue;
                     };
-                    let PubGrubPackage::Package(dependency_package, None, _) = dependency_package
+                    let PubGrubPackage::Package(dependency_package, _, _) = dependency_package
                     else {
                         continue;
                     };
 
+                    // For extras, we include a dependency between the extra and the base package.
+                    if self_package == dependency_package {
+                        continue;
+                    }
+
                     if self_version.contains(version) {
                         let self_index = &inverse[self_package];
                         let dependency_index = &inverse[dependency_package];

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

@@ -607,6 +607,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                     self.visit_package(package, priorities, request_sink)?;
                 }
 
+                // Add a dependency on each editable.
                 for (editable, metadata) in self.editables.values() {
                     constraints.insert(
                         PubGrubPackage::Package(
@@ -616,6 +617,16 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                         ),
                         Range::singleton(metadata.version.clone()),
                     );
+                    for extra in &editable.extras {
+                        constraints.insert(
+                            PubGrubPackage::Package(
+                                metadata.name.clone(),
+                                Some(extra.clone()),
+                                Some(editable.url().clone()),
+                            ),
+                            Range::singleton(metadata.version.clone()),
+                        );
+                    }
                 }
 
                 Ok(Dependencies::Available(constraints.into()))

crates/puffin-resolver/tests/resolver.rs

@@ -185,6 +185,7 @@ async fn black_colorama() -> Result<()> {
     click==8.1.7
         # via black
     colorama==0.4.6
+        # via black
     mypy-extensions==1.0.0
         # via black
     packaging==23.2

crates/puffin/src/commands/pip_compile.rs

@@ -220,8 +220,8 @@ pub(crate) async fn pip_compile(
         let editables: Vec<LocalEditable> = editables
             .into_iter()
             .map(|editable| {
-                let EditableRequirement { path, url } = editable;
-                Ok(LocalEditable { url, path })
+                let EditableRequirement { url, extras, path } = editable;
+                Ok(LocalEditable { url, path, extras })
             })
             .collect::<Result<_>>()?;
 

crates/puffin/src/commands/pip_install.rs

@@ -323,10 +323,11 @@ async fn build_editables(
     let editables: Vec<LocalEditable> = editables
         .iter()
         .map(|editable| {
-            let EditableRequirement { path, url } = editable;
+            let EditableRequirement { url, extras, path } = editable;
             Ok(LocalEditable {
-                path: path.clone(),
                 url: url.clone(),
+                extras: extras.clone(),
+                path: path.clone(),
             })
         })
         .collect::<Result<_>>()?;

crates/puffin/src/commands/pip_sync.rs

@@ -431,10 +431,11 @@ async fn resolve_editables(
         let local_editables: Vec<LocalEditable> = uninstalled
             .iter()
             .map(|editable| {
-                let EditableRequirement { path, url } = editable;
+                let EditableRequirement { url, path, extras } = editable;
                 Ok(LocalEditable {
-                    path: path.clone(),
                     url: url.clone(),
+                    path: path.clone(),
+                    extras: extras.clone(),
                 })
             })
             .collect::<Result<_>>()?;

crates/puffin/tests/pip_compile.rs

@@ -2767,7 +2767,7 @@ fn compile_editable() -> Result<()> {
     requirements_in.write_str(indoc! {r"
         -e ../../scripts/editable-installs/poetry_editable
         -e ${PROJECT_ROOT}/../../scripts/editable-installs/maturin_editable
-        -e file://../../scripts/editable-installs/black_editable
+        -e file://../../scripts/editable-installs/black_editable[d]
         boltons # normal dependency for comparison
         "
     })?;
@@ -2794,16 +2794,34 @@ fn compile_editable() -> Result<()> {
         ----- stdout -----
         # This file was autogenerated by Puffin v[VERSION] via the following command:
         #    puffin pip compile requirements.in --cache-dir [CACHE_DIR]
+        aiohttp==3.9.0
+            # via black
+        aiosignal==1.3.1
+            # via aiohttp
+        attrs==23.1.0
+            # via aiohttp
         -e file://../../scripts/editable-installs/black_editable
         boltons==23.1.1
+        frozenlist==1.4.0
+            # via
+            #   aiohttp
+            #   aiosignal
+        idna==3.4
+            # via yarl
         -e ${PROJECT_ROOT}/../../scripts/editable-installs/maturin_editable
+        multidict==6.0.4
+            # via
+            #   aiohttp
+            #   yarl
         numpy==1.26.2
             # via poetry-editable
         -e ../../scripts/editable-installs/poetry_editable
+        yarl==1.9.2
+            # via aiohttp
 
         ----- stderr -----
         Built 3 editables in [TIME]
-        Resolved 5 packages in [TIME]
+        Resolved 12 packages in [TIME]
         "###);
     });
 

crates/puffin/tests/pip_sync.rs

@@ -2914,13 +2914,12 @@ fn sync_editable_and_registry() -> Result<()> {
         Resolved 1 package in [TIME]
         Downloaded 1 package in [TIME]
         Installed 1 package in [TIME]
-         + black==24.1a1
+         + black==24.1.0
         warning: The package `black` requires `click >=8.0.0`, but it's not installed.
         warning: The package `black` requires `mypy-extensions >=0.4.3`, but it's not installed.
         warning: The package `black` requires `packaging >=22.0`, but it's not installed.
         warning: The package `black` requires `pathspec >=0.9.0`, but it's not installed.
         warning: The package `black` requires `platformdirs >=2`, but it's not installed.
-        warning: The package `black` requires `aiohttp >=3.7.4 ; sys_platform != 'win32' or (implementation_name != 'pypy' and extra == 'd')`, but it's not installed.
         "###);
     });
 
@@ -2961,7 +2960,7 @@ fn sync_editable_and_registry() -> Result<()> {
         Built 1 editable in [TIME]
         Uninstalled 1 package in [TIME]
         Installed 1 package in [TIME]
-         - black==24.1a1
+         - black==24.1.0
          + black==0.1.0 (from file://[WORKSPACE_DIR]/scripts/editable-installs/black_editable)
         "###);
     });

crates/requirements-txt/src/lib.rs

@@ -36,7 +36,6 @@
 
 use std::fmt::{Display, Formatter};
 use std::io;
-use std::io::Error;
 use std::path::{Path, PathBuf};
 
 use fs_err as fs;
@@ -45,8 +44,9 @@ use tracing::warn;
 use unscanny::{Pattern, Scanner};
 use url::Url;
 
-use pep508_rs::{split_scheme, Pep508Error, Pep508ErrorSource, Requirement, VerbatimUrl};
+use pep508_rs::{split_scheme, Extras, Pep508Error, Pep508ErrorSource, Requirement, VerbatimUrl};
 use puffin_fs::NormalizedDisplay;
+use puffin_normalize::ExtraName;
 
 /// We emit one of those for each requirements.txt entry
 enum RequirementsTxtStatement {
@@ -71,6 +71,7 @@ enum RequirementsTxtStatement {
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct EditableRequirement {
     pub url: VerbatimUrl,
+    pub extras: Vec<ExtraName>,
     pub path: PathBuf,
 }
 
@@ -98,8 +99,40 @@ impl EditableRequirement {
         given: &str,
         working_dir: impl AsRef<Path>,
     ) -> Result<EditableRequirement, RequirementsTxtParserError> {
+        // Identify the extras.
+        let (requirement, extras) = if let Some((requirement, extras)) = Self::split_extras(given) {
+            let extras = Extras::parse(extras).map_err(|err| {
+                // Map from error on the extras to error on the whole requirement.
+                let err = Pep508Error {
+                    message: err.message,
+                    start: requirement.len() + err.start,
+                    len: err.len,
+                    input: given.to_string(),
+                };
+                match err.message {
+                    Pep508ErrorSource::String(_) | Pep508ErrorSource::UrlError(_) => {
+                        RequirementsTxtParserError::Pep508 {
+                            start: err.start,
+                            end: err.start + err.len,
+                            source: err,
+                        }
+                    }
+                    Pep508ErrorSource::UnsupportedRequirement(_) => {
+                        RequirementsTxtParserError::UnsupportedRequirement {
+                            start: err.start,
+                            end: err.start + err.len,
+                            source: err,
+                        }
+                    }
+                }
+            })?;
+            (requirement, extras.into_vec())
+        } else {
+            (given, vec![])
+        };
+
         // Create a `VerbatimUrl` to represent the editable requirement.
-        let url = if let Some((scheme, path)) = split_scheme(given) {
+        let url = if let Some((scheme, path)) = split_scheme(requirement) {
             if scheme == "file" {
                 if let Some(path) = path.strip_prefix("//") {
                     // Ex) `file:///home/ferris/project/scripts/...`
@@ -111,23 +144,45 @@ impl EditableRequirement {
             } else {
                 // Ex) `https://...`
                 return Err(RequirementsTxtParserError::UnsupportedUrl(
-                    given.to_string(),
+                    requirement.to_string(),
                 ));
             }
         } else {
             // Ex) `../editable/`
-            VerbatimUrl::from_path(given, working_dir.as_ref())
+            VerbatimUrl::from_path(requirement, working_dir.as_ref())
         };
 
         // Create a `PathBuf`.
-        let path = url
-            .to_file_path()
-            .map_err(|()| RequirementsTxtParserError::InvalidEditablePath(given.to_string()))?;
+        let path = url.to_file_path().map_err(|()| {
+            RequirementsTxtParserError::InvalidEditablePath(requirement.to_string())
+        })?;
 
         // Add the verbatim representation of the URL to the `VerbatimUrl`.
-        let url = url.with_given(given.to_string());
+        let url = url.with_given(requirement.to_string());
 
-        Ok(EditableRequirement { url, path })
+        Ok(EditableRequirement { url, extras, path })
+    }
+
+    /// Identify the extras in an editable URL (e.g., `../editable[dev]`).
+    ///
+    /// Pip uses `m = re.match(r'^(.+)(\[[^]]+])$', path)`. Our strategy is:
+    /// - If the string ends with a closing bracket (`]`)...
+    /// - Iterate backwards until you find the open bracket (`[`)...
+    /// - But abort if you find another closing bracket (`]`) first.
+    pub fn split_extras(given: &str) -> Option<(&str, &str)> {
+        let mut chars = given.char_indices().rev();
+
+        // If the string ends with a closing bracket (`]`)...
+        if !matches!(chars.next(), Some((_, ']'))) {
+            return None;
+        }
+
+        // Iterate backwards until you find the open bracket (`[`)...
+        let (index, _) = chars
+            .take_while(|(_, c)| *c != ']')
+            .find(|(_, c)| *c == '[')?;
+
+        Some(given.split_at(index))
     }
 }
 
@@ -308,7 +363,8 @@ fn parse_entry(
         }
     } else if s.eat_if("-e") {
         let path_or_url = parse_value(s, |c: char| !['\n', '\r'].contains(&c))?;
-        let editable_requirement = EditableRequirement::parse(path_or_url, working_dir)?;
+        let editable_requirement = EditableRequirement::parse(path_or_url, working_dir)
+            .map_err(|err| err.with_offset(start))?;
         RequirementsTxtStatement::EditableRequirement(editable_requirement)
     } else if s.at(char::is_ascii_alphanumeric) {
         let (requirement, hashes) = parse_requirement_and_hashes(s, content, working_dir)?;
@@ -407,16 +463,13 @@ fn parse_requirement_and_hashes(
     let requirement =
         Requirement::parse(&content[start..end], Some(working_dir)).map_err(|err| {
             match err.message {
-                Pep508ErrorSource::String(_) => RequirementsTxtParserError::Pep508 {
-                    source: err,
-                    start,
-                    end,
-                },
-                Pep508ErrorSource::UrlError(_) => RequirementsTxtParserError::Pep508 {
-                    source: err,
-                    start,
-                    end,
-                },
+                Pep508ErrorSource::String(_) | Pep508ErrorSource::UrlError(_) => {
+                    RequirementsTxtParserError::Pep508 {
+                        source: err,
+                        start,
+                        end,
+                    }
+                }
                 Pep508ErrorSource::UnsupportedRequirement(_) => {
                     RequirementsTxtParserError::UnsupportedRequirement {
                         source: err,
@@ -515,6 +568,49 @@ pub enum RequirementsTxtParserError {
     },
 }
 
+impl RequirementsTxtParserError {
+    /// Add a fixed offset to the location of the error.
+    #[must_use]
+    fn with_offset(self, offset: usize) -> Self {
+        match self {
+            RequirementsTxtParserError::IO(err) => RequirementsTxtParserError::IO(err),
+            RequirementsTxtParserError::InvalidEditablePath(given) => {
+                RequirementsTxtParserError::InvalidEditablePath(given)
+            }
+            RequirementsTxtParserError::UnsupportedUrl(url) => {
+                RequirementsTxtParserError::UnsupportedUrl(url)
+            }
+            RequirementsTxtParserError::Parser { message, location } => {
+                RequirementsTxtParserError::Parser {
+                    message,
+                    location: location + offset,
+                }
+            }
+            RequirementsTxtParserError::UnsupportedRequirement { source, start, end } => {
+                RequirementsTxtParserError::UnsupportedRequirement {
+                    source,
+                    start: start + offset,
+                    end: end + offset,
+                }
+            }
+            RequirementsTxtParserError::Pep508 { source, start, end } => {
+                RequirementsTxtParserError::Pep508 {
+                    source,
+                    start: start + offset,
+                    end: end + offset,
+                }
+            }
+            RequirementsTxtParserError::Subfile { source, start, end } => {
+                RequirementsTxtParserError::Subfile {
+                    source,
+                    start: start + offset,
+                    end: end + offset,
+                }
+            }
+        }
+    }
+}
+
 impl Display for RequirementsTxtParserError {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -614,7 +710,7 @@ impl std::error::Error for RequirementsTxtFileError {
 }
 
 impl From<io::Error> for RequirementsTxtParserError {
-    fn from(err: Error) -> Self {
+    fn from(err: io::Error) -> Self {
         Self::IO(err)
     }
 }
@@ -631,7 +727,7 @@ mod test {
     use tempfile::tempdir;
     use test_case::test_case;
 
-    use crate::RequirementsTxt;
+    use crate::{EditableRequirement, RequirementsTxt};
 
     fn workspace_test_data_dir() -> PathBuf {
         PathBuf::from("./test-data")
@@ -774,4 +870,52 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn invalid_editable_extra() -> Result<()> {
+        let temp_dir = assert_fs::TempDir::new()?;
+        let requirements_txt = temp_dir.child("requirements.txt");
+        requirements_txt.write_str(indoc! {"
+            -e black[,abcdef]
+        "})?;
+
+        let error = RequirementsTxt::parse(requirements_txt.path(), temp_dir.path()).unwrap_err();
+        let errors = anyhow::Error::new(error)
+            .chain()
+            .map(|err| err.to_string().replace('\\', "/"))
+            .join("\n");
+
+        insta::with_settings!({
+            filters => vec![(requirements_txt.path().to_str().unwrap(), "<REQUIREMENTS_TXT>")]
+        }, {
+            insta::assert_display_snapshot!(errors, @r###"
+            Couldn't parse requirement in `<REQUIREMENTS_TXT>` at position 6
+            Expected an alphanumeric character starting the extra name, found ','
+            black[,abcdef]
+                  ^
+            "###);
+        });
+
+        Ok(())
+    }
+
+    #[test]
+    fn editable_extra() {
+        assert_eq!(
+            EditableRequirement::split_extras("../editable[dev]"),
+            Some(("../editable", "[dev]"))
+        );
+        assert_eq!(
+            EditableRequirement::split_extras("../editable[dev]more[extra]"),
+            Some(("../editable[dev]more", "[extra]"))
+        );
+        assert_eq!(
+            EditableRequirement::split_extras("../editable[[dev]]"),
+            None
+        );
+        assert_eq!(
+            EditableRequirement::split_extras("../editable[[dev]"),
+            Some(("../editable[", "[dev]"))
+        );
+    }
 }