can use custom pypi indexes

docs/recipe_generation.md

@@ -15,6 +15,23 @@ rattler-build generate-recipe pypi jinja2
 
 This will generate a recipe for the `jinja2` package from PyPI and print it to the console. To turn it into a recipe, you can either pipe the stdout to a file or use the `-w` flag. The `-w` flag will create a new folder with the recipe in it.
 
+The PyPI recipe generation supports additional flags:
+
+- `-w/--write` write the recipe to a folder
+- `-m/--use-mapping` use the conda-forge PyPI name mapping (defaults to true)
+- `-t/--tree` generate recipes for all dependencies
+- `--pypi-index-url` specify one or more PyPI index URLs to use for recipe generation (comma-separated)
+
+The `--pypi-index-url` option allows you to use alternative PyPI mirrors or private PyPI repositories. You can specify multiple URLs, and the system will try each in order until one succeeds. This is especially useful for organizations with private packages or in environments with limited internet access. You can also set the `RATTLER_PYPI_INDEX_URL` environment variable.
+
+```sh
+# Use a custom PyPI index
+rattler-build generate-recipe pypi --pypi-index-url https://my-custom-pypi.example.com/pypi my-package
+
+# Use multiple PyPI indexes (will try each in order)
+rattler-build generate-recipe pypi --pypi-index-url https://my-custom-pypi.example.com/pypi,https://pypi.org/pypi my-package
+```
+
 The generated recipe for `jinja2` will look something like:
 
 ```yaml title="recipe.yaml"

docs/reference/cli.md

@@ -644,6 +644,14 @@ Generate a recipe for a Python package from PyPI
 	Whether to generate recipes for all dependencies
 
 
+- `--pypi-index-url <PYPI_INDEX_URLS>`
+
+	Specify the PyPI index URL(s) to use for recipe generation. Multiple URLs can be specified as a comma-separated list. The system will try each URL in order until one succeeds.
+
+	- Default value: `https://pypi.org/pypi`
+	- Environment variable: `RATTLER_PYPI_INDEX_URL`
+
+
 
 
 

src/recipe_generator/pypi.rs

@@ -36,6 +36,16 @@ pub struct PyPIOpts {
     /// Whether to generate recipes for all dependencies
     #[arg(short, long)]
     pub tree: bool,
+
+    /// Specify the PyPI index URL(s) to use for recipe generation
+    #[arg(
+        long = "pypi-index-url",
+        env = "RATTLER_PYPI_INDEX_URL",
+        default_value = "https://pypi.org/pypi",
+        value_delimiter = ',',
+        help = "Specify the PyPI index URL(s) to use for recipe generation"
+    )]
+    pub pypi_index_urls: Vec<String>,
 }
 
 #[derive(Deserialize, Clone, Debug, Default)]
@@ -254,53 +264,101 @@ async fn fetch_pypi_metadata(
     opts: &PyPIOpts,
     client: &reqwest::Client,
 ) -> miette::Result<PyPiMetadata> {
-    let (info, urls) = if let Some(version) = &opts.version {
-        let url = format!("https://pypi.org/pypi/{}/{}/json", opts.package, version);
-        let release: PyPrReleaseResponse = client
-            .get(&url)
-            .send()
-            .await
-            .into_diagnostic()?
-            .json()
-            .await
-            .into_diagnostic()?;
-        (release.info, release.urls)
-    } else {
-        let url = format!("https://pypi.org/pypi/{}/json", opts.package);
-        let response: PyPiResponse = client
-            .get(&url)
-            .send()
-            .await
-            .into_diagnostic()?
-            .json()
-            .await
-            .into_diagnostic()?;
-
-        // Get the latest release
-        let urls = response
-            .releases
-            .get(&response.info.version)
-            .ok_or_else(|| miette::miette!("No source distribution found"))?;
-        (response.info, urls.clone())
-    };
+    // Try each PyPI index URL in sequence until one works
+    let mut errors = Vec::new();
 
-    let release = urls
-        .iter()
-        .find(|r| r.filename.ends_with(".tar.gz"))
-        .ok_or_else(|| miette::miette!("No source distribution found"))?
-        .clone();
+    for base_url in &opts.pypi_index_urls {
+        // Make sure URL ends with a slash if it doesn't already
+        let base_url = if base_url.ends_with('/') {
+            base_url.to_string()
+        } else {
+            format!("{}/", base_url)
+        };
 
-    let wheel_url = urls
-        .iter()
-        .find(|r| r.filename.ends_with(".whl"))
-        .map(|r| r.url.clone());
-
-    Ok(PyPiMetadata {
-        info,
-        urls,
-        release,
-        wheel_url,
-    })
+        let result: Result<(PyPiInfo, Vec<PyPiRelease>), miette::Error> = async {
+            if let Some(version) = &opts.version {
+                let url = format!("{}{}/{}/json", base_url, opts.package, version);
+                let response = client
+                    .get(&url)
+                    .send()
+                    .await
+                    .into_diagnostic()
+                    .context(format!("Failed to fetch from {}", url))?;
+
+                if !response.status().is_success() {
+                    return Err(miette::miette!(
+                        "Server returned status code: {}",
+                        response.status()
+                    ));
+                }
+
+                let release: PyPrReleaseResponse = response.json().await.into_diagnostic()?;
+                Ok((release.info, release.urls))
+            } else {
+                let url = format!("{}{}/json", base_url, opts.package);
+                let response = client
+                    .get(&url)
+                    .send()
+                    .await
+                    .into_diagnostic()
+                    .context(format!("Failed to fetch from {}", url))?;
+
+                if !response.status().is_success() {
+                    return Err(miette::miette!(
+                        "Server returned status code: {}",
+                        response.status()
+                    ));
+                }
+
+                let response: PyPiResponse = response.json().await.into_diagnostic()?;
+
+                // Get the latest release
+                let urls = response
+                    .releases
+                    .get(&response.info.version)
+                    .ok_or_else(|| miette::miette!("No source distribution found"))?;
+                Ok((response.info, urls.clone()))
+            }
+        }
+        .await;
+
+        match result {
+            Ok((info, urls)) => {
+                // Found a working PyPI index, use it
+                eprintln!("Successfully fetched metadata from {}", base_url);
+
+                let release = urls
+                    .iter()
+                    .find(|r| r.filename.ends_with(".tar.gz"))
+                    .ok_or_else(|| miette::miette!("No source distribution found in {}", base_url))?
+                    .clone();
+
+                let wheel_url = urls
+                    .iter()
+                    .find(|r| r.filename.ends_with(".whl"))
+                    .map(|r| r.url.clone());
+
+                return Ok(PyPiMetadata {
+                    info,
+                    urls,
+                    release,
+                    wheel_url,
+                });
+            }
+            Err(err) => {
+                // Remember the error and try the next URL
+                eprintln!("Failed to fetch from {}: {}", base_url, err);
+                errors.push(format!("{}: {}", base_url, err));
+            }
+        }
+    }
+
+    // If we get here, all URLs failed
+    let error_message = format!(
+        "Failed to fetch metadata from all provided PyPI URLs:\n- {}",
+        errors.join("\n- ")
+    );
+    Err(miette::miette!(error_message))
 }
 
 async fn map_requirement(
@@ -333,11 +391,18 @@ pub async fn create_recipe(
     recipe.package.name = metadata.info.name.to_lowercase();
     recipe.package.version = "${{ version }}".to_string();
 
-    // replace URL with the shorter version that does not contain the hash
-    let release_url = if metadata
-        .release
-        .url
-        .starts_with("https://files.pythonhosted.org/")
+    // Check if we're using the standard PyPI
+    let is_default_pypi = opts
+        .pypi_index_urls
+        .iter()
+        .any(|url| url.starts_with("https://pypi.org"));
+
+    // replace URL with the shorter version that does not contain the hash if using the standard PyPI
+    let release_url = if is_default_pypi
+        && metadata
+            .release
+            .url
+            .starts_with("https://files.pythonhosted.org/")
     {
         let simple_url = format!(
             "https://pypi.org/packages/source/{}/{}/{}-{}.tar.gz",
@@ -467,7 +532,11 @@ pub async fn generate_pypi_recipe(opts: &PyPIOpts) -> miette::Result<()> {
             if !PathBuf::from(dep).exists() {
                 let opts = PyPIOpts {
                     package: dep.to_string(),
-                    ..opts.clone()
+                    version: None,
+                    write: opts.write,
+                    use_mapping: opts.use_mapping,
+                    tree: false, // Don't recursively generate trees
+                    pypi_index_urls: opts.pypi_index_urls.clone(),
                 };
                 generate_pypi_recipe(&opts).await?;
             }
@@ -490,6 +559,7 @@ mod tests {
             write: false,
             use_mapping: true,
             tree: false,
+            pypi_index_urls: vec!["https://pypi.org/pypi".to_string()],
         };
 
         let client = reqwest::Client::new();
@@ -507,6 +577,7 @@ mod tests {
             write: false,
             use_mapping: true,
             tree: false,
+            pypi_index_urls: vec!["https://pypi.org/pypi".to_string()],
         };
 
         let client = reqwest::Client::new();
@@ -515,4 +586,95 @@ mod tests {
 
         assert_yaml_snapshot!(recipe);
     }
+
+    #[tokio::test]
+    async fn test_multiple_pypi_index_urls() {
+        // Test with multiple PyPI index URLs, starting with an invalid one
+        let opts = PyPIOpts {
+            package: "requests".into(),
+            version: Some("2.31.0".into()),
+            write: false,
+            use_mapping: true,
+            tree: false,
+            pypi_index_urls: vec![
+                "https://invalid-pypi-url.example.com/".to_string(),
+                "https://pypi.org/pypi".to_string(),
+            ],
+        };
+
+        let client = reqwest::Client::new();
+        let metadata = fetch_pypi_metadata(&opts, &client).await.unwrap();
+
+        // Verify that the valid URL was used (by checking that metadata was found)
+        assert_eq!(metadata.info.name.to_lowercase(), "requests");
+        assert_eq!(metadata.info.version, "2.31.0");
+    }
+
+    #[test]
+    fn test_format_requirement() {
+        // Test basic requirement formatting
+        assert_eq!(format_requirement("numpy>=1.20.0"), "numpy >=1.20.0");
+
+        // Test requirement with marker
+        assert_eq!(
+            format_requirement("importlib-metadata>=3.6.0;python_version<\"3.10\""),
+            "importlib-metadata >=3.6.0 ;MARKER; python_version<\"3.10\""
+        );
+    }
+
+    #[test]
+    fn test_post_process_markers() {
+        let input = "dependencies:\n- numpy >=1.20.0\n- importlib-metadata >=3.6.0 ;MARKER; python_version<\"3.10\"\n- packaging";
+        let expected = "dependencies:\n- numpy >=1.20.0\n# - importlib-metadata >=3.6.0 # python_version<\"3.10\"\n- packaging";
+
+        assert_eq!(post_process_markers(input.to_string()), expected);
+    }
+
+    #[tokio::test]
+    async fn test_url_simplification() {
+        let client = reqwest::Client::new();
+
+        // Create a PyPI metadata with a PyPI-hosted URL
+        let metadata = PyPiMetadata {
+            info: PyPiInfo {
+                name: "requests".to_string(),
+                version: "2.31.0".to_string(),
+                summary: None,
+                description: None,
+                home_page: None,
+                license: None,
+                requires_dist: None,
+                project_urls: None,
+                requires_python: None,
+            },
+            urls: vec![],
+            release: PyPiRelease {
+                filename: "requests-2.31.0.tar.gz".to_string(),
+                url: "https://files.pythonhosted.org/packages/source/r/requests/requests-2.31.0.tar.gz".to_string(),
+                digests: {
+                    let mut map = HashMap::new();
+                    map.insert("sha256".to_string(), "dummy_hash".to_string());
+                    map
+                },
+            },
+            wheel_url: None,
+        };
+
+        // Create options with default PyPI
+        let opts = PyPIOpts {
+            package: "requests".into(),
+            version: Some("2.31.0".into()),
+            write: false,
+            use_mapping: true,
+            tree: false,
+            pypi_index_urls: vec!["https://pypi.org/pypi".to_string()],
+        };
+
+        let recipe = create_recipe(&opts, &metadata, &client).await.unwrap();
+
+        // Check if the URL was simplified
+        assert!(recipe.source[0]
+            .url
+            .contains("https://pypi.org/packages/source/"));
+    }
 }