1use std::borrow::Cow;
2use std::fmt::{Display, Formatter};
3use std::ops::Deref;
4use std::path::Path;
5use std::str::FromStr;
6use std::sync::{Arc, LazyLock, RwLock};
7
8use itertools::Either;
9use rustc_hash::{FxHashMap, FxHashSet};
10use thiserror::Error;
11use tracing::trace;
12use url::{ParseError, Url};
13use uv_auth::RealmRef;
14use uv_cache_key::CanonicalUrl;
15use uv_pep508::{Scheme, VerbatimUrl, VerbatimUrlError, split_scheme};
16use uv_redacted::DisplaySafeUrl;
17use uv_warnings::warn_user;
18
19use crate::{Index, IndexStatusCodeStrategy, Verbatim};
20
21static PYPI_URL: LazyLock<DisplaySafeUrl> =
22 LazyLock::new(|| DisplaySafeUrl::parse("https://pypi.org/simple").unwrap());
23
24static DEFAULT_INDEX: LazyLock<Index> = LazyLock::new(|| {
25 Index::from_index_url(IndexUrl::Pypi(Arc::new(VerbatimUrl::from_url(
26 PYPI_URL.clone(),
27 ))))
28});
29
30#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd)]
32pub enum IndexUrl {
33 Pypi(Arc<VerbatimUrl>),
34 Url(Arc<VerbatimUrl>),
35 Path(Arc<VerbatimUrl>),
36}
37
38impl IndexUrl {
39 pub fn parse(path: &str, root_dir: Option<&Path>) -> Result<Self, IndexUrlError> {
44 let url = VerbatimUrl::from_url_or_path(path, root_dir)?;
45 Ok(Self::from(url))
46 }
47
48 pub fn root(&self) -> Option<DisplaySafeUrl> {
53 let mut segments = self.url().path_segments()?;
54 let last = match segments.next_back()? {
55 "" => segments.next_back()?,
57 segment => segment,
58 };
59
60 if !(last.eq_ignore_ascii_case("simple") || last.eq_ignore_ascii_case("+simple")) {
62 return None;
63 }
64
65 let mut url = self.url().clone();
66 url.path_segments_mut().ok()?.pop_if_empty().pop();
67 Some(url)
68 }
69}
70
71#[cfg(feature = "schemars")]
72impl schemars::JsonSchema for IndexUrl {
73 fn schema_name() -> Cow<'static, str> {
74 Cow::Borrowed("IndexUrl")
75 }
76
77 fn json_schema(_generator: &mut schemars::generate::SchemaGenerator) -> schemars::Schema {
78 schemars::json_schema!({
79 "type": "string",
80 "description": "The URL of an index to use for fetching packages (e.g., `https://pypi.org/simple`), or a local path."
81 })
82 }
83}
84
85impl IndexUrl {
86 #[inline]
87 fn inner(&self) -> &VerbatimUrl {
88 match self {
89 Self::Pypi(url) | Self::Url(url) | Self::Path(url) => url,
90 }
91 }
92
93 pub fn url(&self) -> &DisplaySafeUrl {
95 self.inner().raw()
96 }
97
98 pub fn into_url(self) -> DisplaySafeUrl {
100 self.inner().to_url()
101 }
102
103 pub fn without_credentials(&self) -> Cow<'_, DisplaySafeUrl> {
105 let url = self.url();
106 if url.username().is_empty() && url.password().is_none() {
107 Cow::Borrowed(url)
108 } else {
109 let mut url = url.clone();
110 let _ = url.set_username("");
111 let _ = url.set_password(None);
112 Cow::Owned(url)
113 }
114 }
115
116 pub fn warn_on_disambiguated_relative_path(&self) {
121 let Self::Path(verbatim_url) = &self else {
122 return;
123 };
124
125 if let Some(path) = verbatim_url.given() {
126 if !is_disambiguated_path(path) {
127 if cfg!(windows) {
128 warn_user!(
129 "Relative paths passed to `--index` or `--default-index` should be disambiguated from index names (use `.\\{path}` or `./{path}`). Support for ambiguous values will be removed in the future"
130 );
131 } else {
132 warn_user!(
133 "Relative paths passed to `--index` or `--default-index` should be disambiguated from index names (use `./{path}`). Support for ambiguous values will be removed in the future"
134 );
135 }
136 }
137 }
138 }
139}
140
141impl Display for IndexUrl {
142 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
143 Display::fmt(self.inner(), f)
144 }
145}
146
147impl Verbatim for IndexUrl {
148 fn verbatim(&self) -> Cow<'_, str> {
149 self.inner().verbatim()
150 }
151}
152
153fn is_disambiguated_path(path: &str) -> bool {
159 if cfg!(windows) {
160 if path.starts_with(".\\") || path.starts_with("..\\") || path.starts_with('/') {
161 return true;
162 }
163 }
164 if path.starts_with("./") || path.starts_with("../") || Path::new(path).is_absolute() {
165 return true;
166 }
167 if let Some((scheme, _)) = split_scheme(path) {
169 return Scheme::parse(scheme).is_some();
170 }
171 false
173}
174
175#[derive(Error, Debug)]
177pub enum IndexUrlError {
178 #[error(transparent)]
179 Io(#[from] std::io::Error),
180 #[error(transparent)]
181 Url(#[from] ParseError),
182 #[error(transparent)]
183 VerbatimUrl(#[from] VerbatimUrlError),
184}
185
186impl FromStr for IndexUrl {
187 type Err = IndexUrlError;
188
189 fn from_str(s: &str) -> Result<Self, Self::Err> {
190 Self::parse(s, None)
191 }
192}
193
194impl serde::ser::Serialize for IndexUrl {
195 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
196 where
197 S: serde::ser::Serializer,
198 {
199 self.inner().without_credentials().serialize(serializer)
200 }
201}
202
203impl<'de> serde::de::Deserialize<'de> for IndexUrl {
204 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
205 where
206 D: serde::de::Deserializer<'de>,
207 {
208 struct Visitor;
209
210 impl serde::de::Visitor<'_> for Visitor {
211 type Value = IndexUrl;
212
213 fn expecting(&self, f: &mut Formatter) -> std::fmt::Result {
214 f.write_str("a string")
215 }
216
217 fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<Self::Value, E> {
218 IndexUrl::from_str(v).map_err(serde::de::Error::custom)
219 }
220 }
221
222 deserializer.deserialize_str(Visitor)
223 }
224}
225
226impl From<VerbatimUrl> for IndexUrl {
227 fn from(url: VerbatimUrl) -> Self {
228 if url.scheme() == "file" {
229 Self::Path(Arc::new(url))
230 } else if *url.raw() == *PYPI_URL {
231 Self::Pypi(Arc::new(url))
232 } else {
233 Self::Url(Arc::new(url))
234 }
235 }
236}
237
238impl From<IndexUrl> for DisplaySafeUrl {
239 fn from(index: IndexUrl) -> Self {
240 index.inner().to_url()
241 }
242}
243
244impl Deref for IndexUrl {
245 type Target = Url;
246
247 fn deref(&self) -> &Self::Target {
248 self.inner()
249 }
250}
251
252#[derive(Default, Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
257#[serde(rename_all = "kebab-case", deny_unknown_fields)]
258pub struct IndexLocations {
259 indexes: Vec<Index>,
260 flat_index: Vec<Index>,
261 no_index: bool,
262}
263
264impl IndexLocations {
265 pub fn new(indexes: Vec<Index>, flat_index: Vec<Index>, no_index: bool) -> Self {
267 Self {
268 indexes,
269 flat_index,
270 no_index,
271 }
272 }
273
274 #[must_use]
281 pub fn combine(self, indexes: Vec<Index>, flat_index: Vec<Index>, no_index: bool) -> Self {
282 Self {
283 indexes: self.indexes.into_iter().chain(indexes).collect(),
284 flat_index: self.flat_index.into_iter().chain(flat_index).collect(),
285 no_index: self.no_index || no_index,
286 }
287 }
288
289 pub fn is_none(&self) -> bool {
292 *self == Self::default()
293 }
294}
295
296fn is_same_index(a: &IndexUrl, b: &IndexUrl) -> bool {
298 RealmRef::from(&**b.url()) == RealmRef::from(&**a.url())
299 && CanonicalUrl::new(a.url()) == CanonicalUrl::new(b.url())
300}
301
302impl<'a> IndexLocations {
303 pub fn default_index(&'a self) -> Option<&'a Index> {
309 if self.no_index {
310 None
311 } else {
312 let mut seen = FxHashSet::default();
313 self.indexes
314 .iter()
315 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
316 .find(|index| index.default)
317 .or_else(|| Some(&DEFAULT_INDEX))
318 }
319 }
320
321 pub fn implicit_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
325 if self.no_index {
326 Either::Left(std::iter::empty())
327 } else {
328 let mut seen = FxHashSet::default();
329 Either::Right(
330 self.indexes
331 .iter()
332 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
333 .filter(|index| !index.default && !index.explicit),
334 )
335 }
336 }
337
338 pub fn indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
347 self.implicit_indexes()
348 .chain(self.default_index())
349 .filter(|index| !index.explicit)
350 }
351
352 pub fn simple_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
356 if self.no_index {
357 Either::Left(std::iter::empty())
358 } else {
359 let mut seen = FxHashSet::default();
360 Either::Right(
361 self.indexes
362 .iter()
363 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name))),
364 )
365 }
366 }
367
368 pub fn flat_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
370 self.flat_index.iter()
371 }
372
373 pub fn no_index(&self) -> bool {
375 self.no_index
376 }
377
378 pub fn index_urls(&'a self) -> IndexUrls {
380 IndexUrls {
381 indexes: self.indexes.clone(),
382 no_index: self.no_index,
383 }
384 }
385
386 pub fn allowed_indexes(&'a self) -> Vec<&'a Index> {
393 if self.no_index {
394 self.flat_index.iter().rev().collect()
395 } else {
396 let mut indexes = vec![];
397
398 let mut seen = FxHashSet::default();
399 let mut default = false;
400 for index in {
401 self.indexes
402 .iter()
403 .chain(self.flat_index.iter())
404 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
405 } {
406 if index.default {
407 if default {
408 continue;
409 }
410 default = true;
411 }
412 indexes.push(index);
413 }
414 if !default {
415 indexes.push(&*DEFAULT_INDEX);
416 }
417
418 indexes.reverse();
419 indexes
420 }
421 }
422
423 pub fn known_indexes(&'a self) -> impl Iterator<Item = &'a Index> {
432 if self.no_index {
433 Either::Left(self.flat_index.iter().rev())
434 } else {
435 Either::Right(
436 std::iter::once(&*DEFAULT_INDEX)
437 .chain(self.flat_index.iter().rev())
438 .chain(self.indexes.iter().rev()),
439 )
440 }
441 }
442
443 pub fn cache_index_credentials(&self) {
445 for index in self.known_indexes() {
446 if let Some(credentials) = index.credentials() {
447 trace!(
448 "Read credentials for index {}",
449 index
450 .name
451 .as_ref()
452 .map(ToString::to_string)
453 .unwrap_or_else(|| index.url.to_string())
454 );
455 if let Some(root_url) = index.root_url() {
456 uv_auth::store_credentials(&root_url, credentials.clone());
457 }
458 uv_auth::store_credentials(index.raw_url(), credentials);
459 }
460 }
461 }
462
463 pub fn simple_api_cache_control_for(&self, url: &IndexUrl) -> Option<&str> {
465 for index in &self.indexes {
466 if is_same_index(index.url(), url) {
467 return index.simple_api_cache_control();
468 }
469 }
470 None
471 }
472
473 pub fn artifact_cache_control_for(&self, url: &IndexUrl) -> Option<&str> {
475 for index in &self.indexes {
476 if is_same_index(index.url(), url) {
477 return index.artifact_cache_control();
478 }
479 }
480 None
481 }
482}
483
484impl From<&IndexLocations> for uv_auth::Indexes {
485 fn from(index_locations: &IndexLocations) -> Self {
486 Self::from_indexes(index_locations.allowed_indexes().into_iter().map(|index| {
487 let mut url = index.url().url().clone();
488 url.set_username("").ok();
489 url.set_password(None).ok();
490 let mut root_url = index.url().root().unwrap_or_else(|| url.clone());
491 root_url.set_username("").ok();
492 root_url.set_password(None).ok();
493 uv_auth::Index {
494 url,
495 root_url,
496 auth_policy: index.authenticate,
497 }
498 }))
499 }
500}
501
502#[derive(Default, Debug, Clone, PartialEq, Eq)]
507pub struct IndexUrls {
508 indexes: Vec<Index>,
509 no_index: bool,
510}
511
512impl<'a> IndexUrls {
513 pub fn from_indexes(indexes: Vec<Index>) -> Self {
514 Self {
515 indexes,
516 no_index: false,
517 }
518 }
519
520 fn default_index(&'a self) -> Option<&'a Index> {
526 if self.no_index {
527 None
528 } else {
529 let mut seen = FxHashSet::default();
530 self.indexes
531 .iter()
532 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
533 .find(|index| index.default)
534 .or_else(|| Some(&DEFAULT_INDEX))
535 }
536 }
537
538 fn implicit_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
542 if self.no_index {
543 Either::Left(std::iter::empty())
544 } else {
545 let mut seen = FxHashSet::default();
546 Either::Right(
547 self.indexes
548 .iter()
549 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
550 .filter(|index| !index.default && !index.explicit),
551 )
552 }
553 }
554
555 pub fn indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
563 let mut seen = FxHashSet::default();
564 self.implicit_indexes()
565 .chain(self.default_index())
566 .filter(|index| !index.explicit)
567 .filter(move |index| seen.insert(index.raw_url())) }
569
570 pub fn defined_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
581 if self.no_index {
582 return Either::Left(std::iter::empty());
583 }
584
585 let mut seen = FxHashSet::default();
586 let (non_default, default) = self
587 .indexes
588 .iter()
589 .filter(move |index| {
590 if let Some(name) = &index.name {
591 seen.insert(name)
592 } else {
593 true
594 }
595 })
596 .partition::<Vec<_>, _>(|index| !index.default);
597
598 Either::Right(non_default.into_iter().chain(default))
599 }
600
601 pub fn no_index(&self) -> bool {
603 self.no_index
604 }
605
606 pub fn status_code_strategy_for(&self, url: &IndexUrl) -> IndexStatusCodeStrategy {
608 for index in &self.indexes {
609 if is_same_index(index.url(), url) {
610 return index.status_code_strategy();
611 }
612 }
613 IndexStatusCodeStrategy::Default
614 }
615
616 pub fn simple_api_cache_control_for(&self, url: &IndexUrl) -> Option<&str> {
618 for index in &self.indexes {
619 if is_same_index(index.url(), url) {
620 return index.simple_api_cache_control();
621 }
622 }
623 None
624 }
625
626 pub fn artifact_cache_control_for(&self, url: &IndexUrl) -> Option<&str> {
628 for index in &self.indexes {
629 if is_same_index(index.url(), url) {
630 return index.artifact_cache_control();
631 }
632 }
633 None
634 }
635}
636
637bitflags::bitflags! {
638 #[derive(Debug, Copy, Clone)]
639 struct Flags: u8 {
640 const NO_RANGE_REQUESTS = 1;
642 const UNAUTHORIZED = 1 << 2;
644 const FORBIDDEN = 1 << 1;
646 }
647}
648
649#[derive(Debug, Default, Clone)]
655pub struct IndexCapabilities(Arc<RwLock<FxHashMap<IndexUrl, Flags>>>);
656
657impl IndexCapabilities {
658 pub fn supports_range_requests(&self, index_url: &IndexUrl) -> bool {
660 !self
661 .0
662 .read()
663 .unwrap()
664 .get(index_url)
665 .is_some_and(|flags| flags.intersects(Flags::NO_RANGE_REQUESTS))
666 }
667
668 pub fn set_no_range_requests(&self, index_url: IndexUrl) {
670 self.0
671 .write()
672 .unwrap()
673 .entry(index_url)
674 .or_insert(Flags::empty())
675 .insert(Flags::NO_RANGE_REQUESTS);
676 }
677
678 pub fn unauthorized(&self, index_url: &IndexUrl) -> bool {
680 self.0
681 .read()
682 .unwrap()
683 .get(index_url)
684 .is_some_and(|flags| flags.intersects(Flags::UNAUTHORIZED))
685 }
686
687 pub fn set_unauthorized(&self, index_url: IndexUrl) {
689 self.0
690 .write()
691 .unwrap()
692 .entry(index_url)
693 .or_insert(Flags::empty())
694 .insert(Flags::UNAUTHORIZED);
695 }
696
697 pub fn forbidden(&self, index_url: &IndexUrl) -> bool {
699 self.0
700 .read()
701 .unwrap()
702 .get(index_url)
703 .is_some_and(|flags| flags.intersects(Flags::FORBIDDEN))
704 }
705
706 pub fn set_forbidden(&self, index_url: IndexUrl) {
708 self.0
709 .write()
710 .unwrap()
711 .entry(index_url)
712 .or_insert(Flags::empty())
713 .insert(Flags::FORBIDDEN);
714 }
715}
716
717#[cfg(test)]
718mod tests {
719 use super::*;
720 use crate::{IndexCacheControl, IndexFormat, IndexName};
721 use uv_small_str::SmallString;
722
723 #[test]
724 fn test_index_url_parse_valid_paths() {
725 assert!(is_disambiguated_path("/absolute/path"));
727 assert!(is_disambiguated_path("./relative/path"));
729 assert!(is_disambiguated_path("../../relative/path"));
730 if cfg!(windows) {
731 assert!(is_disambiguated_path("C:/absolute/path"));
733 assert!(is_disambiguated_path(".\\relative\\path"));
735 assert!(is_disambiguated_path("..\\..\\relative\\path"));
736 }
737 }
738
739 #[test]
740 fn test_index_url_parse_ambiguous_paths() {
741 assert!(!is_disambiguated_path("index"));
743 assert!(!is_disambiguated_path("relative/path"));
745 }
746
747 #[test]
748 fn test_index_url_parse_with_schemes() {
749 assert!(is_disambiguated_path("file:///absolute/path"));
750 assert!(is_disambiguated_path("https://registry.com/simple/"));
751 assert!(is_disambiguated_path(
752 "git+https://github.com/example/repo.git"
753 ));
754 }
755
756 #[test]
757 fn test_cache_control_lookup() {
758 use std::str::FromStr;
759
760 use uv_small_str::SmallString;
761
762 use crate::IndexFormat;
763 use crate::index_name::IndexName;
764
765 let indexes = vec![
766 Index {
767 name: Some(IndexName::from_str("index1").unwrap()),
768 url: IndexUrl::from_str("https://index1.example.com/simple").unwrap(),
769 cache_control: Some(crate::IndexCacheControl {
770 api: Some(SmallString::from("max-age=300")),
771 files: Some(SmallString::from("max-age=1800")),
772 }),
773 explicit: false,
774 default: false,
775 origin: None,
776 format: IndexFormat::Simple,
777 publish_url: None,
778 authenticate: uv_auth::AuthPolicy::default(),
779 ignore_error_codes: None,
780 },
781 Index {
782 name: Some(IndexName::from_str("index2").unwrap()),
783 url: IndexUrl::from_str("https://index2.example.com/simple").unwrap(),
784 cache_control: None,
785 explicit: false,
786 default: false,
787 origin: None,
788 format: IndexFormat::Simple,
789 publish_url: None,
790 authenticate: uv_auth::AuthPolicy::default(),
791 ignore_error_codes: None,
792 },
793 ];
794
795 let index_urls = IndexUrls::from_indexes(indexes);
796
797 let url1 = IndexUrl::from_str("https://index1.example.com/simple").unwrap();
798 assert_eq!(
799 index_urls.simple_api_cache_control_for(&url1),
800 Some("max-age=300")
801 );
802 assert_eq!(
803 index_urls.artifact_cache_control_for(&url1),
804 Some("max-age=1800")
805 );
806
807 let url2 = IndexUrl::from_str("https://index2.example.com/simple").unwrap();
808 assert_eq!(index_urls.simple_api_cache_control_for(&url2), None);
809 assert_eq!(index_urls.artifact_cache_control_for(&url2), None);
810
811 let url3 = IndexUrl::from_str("https://index3.example.com/simple").unwrap();
812 assert_eq!(index_urls.simple_api_cache_control_for(&url3), None);
813 assert_eq!(index_urls.artifact_cache_control_for(&url3), None);
814 }
815
816 #[test]
817 fn test_pytorch_default_cache_control() {
818 let indexes = vec![Index {
820 name: Some(IndexName::from_str("pytorch").unwrap()),
821 url: IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap(),
822 cache_control: None, explicit: false,
824 default: false,
825 origin: None,
826 format: IndexFormat::Simple,
827 publish_url: None,
828 authenticate: uv_auth::AuthPolicy::default(),
829 ignore_error_codes: None,
830 }];
831
832 let index_urls = IndexUrls::from_indexes(indexes.clone());
833 let index_locations = IndexLocations::new(indexes, Vec::new(), false);
834
835 let pytorch_url = IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap();
836
837 assert_eq!(index_urls.simple_api_cache_control_for(&pytorch_url), None);
839 assert_eq!(
840 index_urls.artifact_cache_control_for(&pytorch_url),
841 Some("max-age=365000000, immutable, public")
842 );
843
844 assert_eq!(
846 index_locations.simple_api_cache_control_for(&pytorch_url),
847 None
848 );
849 assert_eq!(
850 index_locations.artifact_cache_control_for(&pytorch_url),
851 Some("max-age=365000000, immutable, public")
852 );
853 }
854
855 #[test]
856 fn test_pytorch_user_override_cache_control() {
857 let indexes = vec![Index {
859 name: Some(IndexName::from_str("pytorch").unwrap()),
860 url: IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap(),
861 cache_control: Some(IndexCacheControl {
862 api: Some(SmallString::from("no-cache")),
863 files: Some(SmallString::from("max-age=3600")),
864 }),
865 explicit: false,
866 default: false,
867 origin: None,
868 format: IndexFormat::Simple,
869 publish_url: None,
870 authenticate: uv_auth::AuthPolicy::default(),
871 ignore_error_codes: None,
872 }];
873
874 let index_urls = IndexUrls::from_indexes(indexes.clone());
875 let index_locations = IndexLocations::new(indexes, Vec::new(), false);
876
877 let pytorch_url = IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap();
878
879 assert_eq!(
881 index_urls.simple_api_cache_control_for(&pytorch_url),
882 Some("no-cache")
883 );
884 assert_eq!(
885 index_urls.artifact_cache_control_for(&pytorch_url),
886 Some("max-age=3600")
887 );
888
889 assert_eq!(
891 index_locations.simple_api_cache_control_for(&pytorch_url),
892 Some("no-cache")
893 );
894 assert_eq!(
895 index_locations.artifact_cache_control_for(&pytorch_url),
896 Some("max-age=3600")
897 );
898 }
899}