From 467dbfbe84f08ed909067521670a73207df8d027 Mon Sep 17 00:00:00 2001
From: figsoda <figsoda@pm.me>
Date: Wed, 15 Feb 2023 18:31:00 -0500
Subject: [PATCH] optimize fetchPypi with `extension` specified

---
 src/fetcher/pypi.rs                           | 36 +++++++++++++------
 .../{pypi.stdout => pypi/basic.stdout}        |  0
 .../fetcher/{pypi.toml => pypi/basic.toml}    |  0
 tests/cmd/fetcher/pypi/extension.stdout       |  6 ++++
 tests/cmd/fetcher/pypi/extension.toml         |  1 +
 5 files changed, 32 insertions(+), 11 deletions(-)
 rename tests/cmd/fetcher/{pypi.stdout => pypi/basic.stdout} (100%)
 rename tests/cmd/fetcher/{pypi.toml => pypi/basic.toml} (100%)
 create mode 100644 tests/cmd/fetcher/pypi/extension.stdout
 create mode 100644 tests/cmd/fetcher/pypi/extension.toml

diff --git a/src/fetcher/pypi.rs b/src/fetcher/pypi.rs
index 27f547c..9166463 100644
--- a/src/fetcher/pypi.rs
+++ b/src/fetcher/pypi.rs
@@ -1,8 +1,6 @@
-use crate::{
-    impl_fetcher,
-    simple::{SimpleFetcher, SimpleUrlFetcher},
-    Url,
-};
+use anyhow::Result;
+
+use crate::{impl_fetcher, prefetch::url_prefetch, simple::SimpleFetcher, Url};
 
 pub struct FetchPypi;
 impl_fetcher!(FetchPypi);
@@ -18,11 +16,27 @@ impl<'a> SimpleFetcher<'a, 1> for FetchPypi {
     }
 }
 
-impl SimpleUrlFetcher<'_, 1> for FetchPypi {
-    const UNPACK: bool = false;
-
-    fn get_url(&self, [pname]: &[&str; 1], version: &str) -> String {
-        let Some(first) = pname.chars().next() else { unreachable!(); };
-        format!("https://pypi.org/packages/source/{first}/{pname}/{pname}-{version}.tar.gz")
+impl FetchPypi {
+    fn fetch(
+        &self,
+        values @ [pname]: &[&str; 1],
+        version: &str,
+        submodules: bool,
+        args: &[(String, String)],
+        args_str: &[(String, String)],
+        nixpkgs: String,
+    ) -> Result<String> {
+        match (args, args_str) {
+            ([], []) => url_prefetch(get_url(pname, version, "tar.gz"), false),
+            ([], [(key, ext)]) if key == "extension" => {
+                url_prefetch(get_url(pname, version, ext), false)
+            }
+            _ => self.fetch_fod(values, version, submodules, args, args_str, nixpkgs),
+        }
     }
 }
+
+fn get_url(pname: &str, version: &str, ext: &str) -> String {
+    let Some(first) = pname.chars().next() else { unreachable!(); };
+    format!("https://pypi.org/packages/source/{first}/{pname}/{pname}-{version}.{ext}")
+}
diff --git a/tests/cmd/fetcher/pypi.stdout b/tests/cmd/fetcher/pypi/basic.stdout
similarity index 100%
rename from tests/cmd/fetcher/pypi.stdout
rename to tests/cmd/fetcher/pypi/basic.stdout
diff --git a/tests/cmd/fetcher/pypi.toml b/tests/cmd/fetcher/pypi/basic.toml
similarity index 100%
rename from tests/cmd/fetcher/pypi.toml
rename to tests/cmd/fetcher/pypi/basic.toml
diff --git a/tests/cmd/fetcher/pypi/extension.stdout b/tests/cmd/fetcher/pypi/extension.stdout
new file mode 100644
index 0000000..675576d
--- /dev/null
+++ b/tests/cmd/fetcher/pypi/extension.stdout
@@ -0,0 +1,6 @@
+fetchPypi {
+  pname = "nltk";
+  version = "3.8";
+  hash = "sha256-dLMIJqN9eNU0JxBbvQN92IAlG+Jp/KZO5TCDikbtVfw=";
+  extension = "zip";
+}
\ No newline at end of file
diff --git a/tests/cmd/fetcher/pypi/extension.toml b/tests/cmd/fetcher/pypi/extension.toml
new file mode 100644
index 0000000..13641bf
--- /dev/null
+++ b/tests/cmd/fetcher/pypi/extension.toml
@@ -0,0 +1 @@
+args = ["https://pypi.org/project/nltk", "3.8", "--arg-str", "extension", "zip"]