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 url::{ParseError, Url};
12use uv_auth::RealmRef;
13use uv_cache_key::CanonicalUrl;
14use uv_pep508::{Scheme, VerbatimUrl, VerbatimUrlError, split_scheme};
15use uv_redacted::DisplaySafeUrl;
16use uv_warnings::warn_user;
17
18use crate::{ExcludeNewerOverride, Index, IndexStatusCodeStrategy, Verbatim};
19
20pub static PYPI_URL: LazyLock<DisplaySafeUrl> =
21 LazyLock::new(|| DisplaySafeUrl::parse("https://pypi.org/simple").unwrap());
22
23static DEFAULT_INDEX: LazyLock<Index> = LazyLock::new(|| {
24 Index::from_index_url(IndexUrl::Pypi(Arc::new(VerbatimUrl::from_url(
25 PYPI_URL.clone(),
26 ))))
27});
28
29#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd)]
31pub enum IndexUrl {
32 Pypi(Arc<VerbatimUrl>),
33 Url(Arc<VerbatimUrl>),
34 Path(Arc<VerbatimUrl>),
35}
36
37impl IndexUrl {
38 pub fn parse(path: &str, root_dir: Option<&Path>) -> Result<Self, IndexUrlError> {
43 let url = VerbatimUrl::from_url_or_path(path, root_dir)?;
44 Ok(Self::from(url))
45 }
46
47 pub fn root(&self) -> Option<DisplaySafeUrl> {
52 let mut segments = self.url().path_segments()?;
53 let last = match segments.next_back()? {
54 "" => segments.next_back()?,
56 segment => segment,
57 };
58
59 if !(last.eq_ignore_ascii_case("simple") || last.eq_ignore_ascii_case("+simple")) {
61 return None;
62 }
63
64 let mut url = self.url().clone();
65 url.path_segments_mut().ok()?.pop_if_empty().pop();
66 Some(url)
67 }
68}
69
70#[cfg(feature = "schemars")]
71impl schemars::JsonSchema for IndexUrl {
72 fn schema_name() -> Cow<'static, str> {
73 Cow::Borrowed("IndexUrl")
74 }
75
76 fn json_schema(_generator: &mut schemars::generate::SchemaGenerator) -> schemars::Schema {
77 schemars::json_schema!({
78 "type": "string",
79 "description": "The URL of an index to use for fetching packages (e.g., `https://pypi.org/simple`), or a local path."
80 })
81 }
82}
83
84impl IndexUrl {
85 #[inline]
86 fn inner(&self) -> &VerbatimUrl {
87 match self {
88 Self::Pypi(url) | Self::Url(url) | Self::Path(url) => url,
89 }
90 }
91
92 pub fn url(&self) -> &DisplaySafeUrl {
94 self.inner().raw()
95 }
96
97 pub fn into_url(self) -> DisplaySafeUrl {
99 self.inner().to_url()
100 }
101
102 pub fn without_credentials(&self) -> Cow<'_, DisplaySafeUrl> {
104 let url = self.url();
105 if url.username().is_empty() && url.password().is_none() {
106 Cow::Borrowed(url)
107 } else {
108 let mut url = url.clone();
109 let _ = url.set_username("");
110 let _ = url.set_password(None);
111 Cow::Owned(url)
112 }
113 }
114
115 pub fn warn_on_disambiguated_relative_path(&self) {
120 let Self::Path(verbatim_url) = &self else {
121 return;
122 };
123
124 if let Some(path) = verbatim_url.given()
125 && !is_disambiguated_path(path)
126 {
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
140impl Display for IndexUrl {
141 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
142 Display::fmt(self.inner(), f)
143 }
144}
145
146impl Verbatim for IndexUrl {
147 fn verbatim(&self) -> Cow<'_, str> {
148 self.inner().verbatim()
149 }
150}
151
152fn is_disambiguated_path(path: &str) -> bool {
158 if cfg!(windows) {
159 if path.starts_with(".\\") || path.starts_with("..\\") || path.starts_with('/') {
160 return true;
161 }
162 }
163 if path.starts_with("./") || path.starts_with("../") || Path::new(path).is_absolute() {
164 return true;
165 }
166 if let Some((scheme, _)) = split_scheme(path) {
168 return Scheme::parse(scheme).is_some();
169 }
170 false
172}
173
174#[derive(Error, Debug)]
176pub enum IndexUrlError {
177 #[error(transparent)]
178 Io(#[from] std::io::Error),
179 #[error(transparent)]
180 Url(#[from] ParseError),
181 #[error(transparent)]
182 VerbatimUrl(#[from] VerbatimUrlError),
183}
184
185impl FromStr for IndexUrl {
186 type Err = IndexUrlError;
187
188 fn from_str(s: &str) -> Result<Self, Self::Err> {
189 Self::parse(s, None)
190 }
191}
192
193impl serde::ser::Serialize for IndexUrl {
194 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
195 where
196 S: serde::ser::Serializer,
197 {
198 self.inner().without_credentials().serialize(serializer)
199 }
200}
201
202impl<'de> serde::de::Deserialize<'de> for IndexUrl {
203 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
204 where
205 D: serde::de::Deserializer<'de>,
206 {
207 struct Visitor;
208
209 impl serde::de::Visitor<'_> for Visitor {
210 type Value = IndexUrl;
211
212 fn expecting(&self, f: &mut Formatter) -> std::fmt::Result {
213 f.write_str("a string")
214 }
215
216 fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<Self::Value, E> {
217 IndexUrl::from_str(v).map_err(serde::de::Error::custom)
218 }
219 }
220
221 deserializer.deserialize_str(Visitor)
222 }
223}
224
225impl From<VerbatimUrl> for IndexUrl {
226 fn from(url: VerbatimUrl) -> Self {
227 if url.scheme() == "file" {
228 Self::Path(Arc::new(url))
229 } else if *url.raw() == *PYPI_URL {
230 Self::Pypi(Arc::new(url))
231 } else {
232 Self::Url(Arc::new(url))
233 }
234 }
235}
236
237impl From<IndexUrl> for DisplaySafeUrl {
238 fn from(index: IndexUrl) -> Self {
239 index.inner().to_url()
240 }
241}
242
243impl Deref for IndexUrl {
244 type Target = Url;
245
246 fn deref(&self) -> &Self::Target {
247 self.inner()
248 }
249}
250
251#[derive(Default, Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
256#[serde(rename_all = "kebab-case", deny_unknown_fields)]
257pub struct IndexLocations {
258 indexes: Vec<Index>,
259 flat_index: Vec<Index>,
260 no_index: bool,
261}
262
263impl IndexLocations {
264 pub fn new(indexes: Vec<Index>, flat_index: Vec<Index>, no_index: bool) -> Self {
266 Self {
267 indexes,
268 flat_index,
269 no_index,
270 }
271 }
272
273 #[must_use]
280 pub fn combine(self, indexes: Vec<Index>, flat_index: Vec<Index>, no_index: bool) -> Self {
281 Self {
282 indexes: self.indexes.into_iter().chain(indexes).collect(),
283 flat_index: self.flat_index.into_iter().chain(flat_index).collect(),
284 no_index: self.no_index || no_index,
285 }
286 }
287
288 pub fn is_none(&self) -> bool {
291 *self == Self::default()
292 }
293}
294
295fn is_same_index(a: &IndexUrl, b: &IndexUrl) -> bool {
297 RealmRef::from(&**b.url()) == RealmRef::from(&**a.url())
298 && CanonicalUrl::new(a.url()) == CanonicalUrl::new(b.url())
299}
300
301impl<'a> IndexLocations {
302 pub fn default_index(&'a self) -> Option<&'a Index> {
308 if self.no_index {
309 None
310 } else {
311 let mut seen = FxHashSet::default();
312 self.indexes
313 .iter()
314 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
315 .find(|index| index.default)
316 .or_else(|| Some(&DEFAULT_INDEX))
317 }
318 }
319
320 pub fn implicit_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
324 if self.no_index {
325 Either::Left(std::iter::empty())
326 } else {
327 let mut seen = FxHashSet::default();
328 Either::Right(
329 self.indexes
330 .iter()
331 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
332 .filter(|index| !index.default && !index.explicit),
333 )
334 }
335 }
336
337 pub fn explicit_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
341 if self.no_index {
342 Either::Left(std::iter::empty())
343 } else {
344 let mut seen = FxHashSet::default();
345 Either::Right(
346 self.indexes
347 .iter()
348 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
349 .filter(|index| index.explicit),
350 )
351 }
352 }
353
354 pub fn indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
363 self.implicit_indexes()
364 .chain(self.default_index())
365 .filter(|index| !index.explicit)
366 }
367
368 pub fn fetch_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
372 let mut seen = FxHashSet::default();
373 self.indexes()
374 .filter(move |index| seen.insert(index.raw_url()))
375 }
376
377 pub fn simple_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
381 if self.no_index {
382 Either::Left(std::iter::empty())
383 } else {
384 let mut seen = FxHashSet::default();
385 Either::Right(
386 self.indexes
387 .iter()
388 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name))),
389 )
390 }
391 }
392
393 pub fn flat_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
395 self.flat_index.iter()
396 }
397
398 pub fn no_index(&self) -> bool {
400 self.no_index
401 }
402
403 pub fn allowed_indexes(&'a self) -> Vec<&'a Index> {
410 if self.no_index {
411 self.flat_index.iter().rev().collect()
412 } else {
413 let mut indexes = vec![];
414
415 let mut seen = FxHashSet::default();
416 let mut default = false;
417 for index in {
418 self.indexes
419 .iter()
420 .chain(self.flat_index.iter())
421 .filter(move |index| index.name.as_ref().is_none_or(|name| seen.insert(name)))
422 } {
423 if index.default {
424 if default {
425 continue;
426 }
427 default = true;
428 }
429 indexes.push(index);
430 }
431 if !default {
432 indexes.push(&*DEFAULT_INDEX);
433 }
434
435 indexes.reverse();
436 indexes
437 }
438 }
439
440 pub fn known_indexes(&'a self) -> impl Iterator<Item = &'a Index> {
449 if self.no_index {
450 Either::Left(self.flat_index.iter().rev())
451 } else {
452 Either::Right(
453 std::iter::once(&*DEFAULT_INDEX)
454 .chain(self.flat_index.iter().rev())
455 .chain(self.indexes.iter().rev()),
456 )
457 }
458 }
459
460 pub fn defined_indexes(&'a self) -> impl Iterator<Item = &'a Index> + 'a {
470 if self.no_index {
471 return Either::Left(std::iter::empty());
472 }
473
474 let mut seen = FxHashSet::default();
475 let (non_default, default) = self
476 .indexes
477 .iter()
478 .filter(move |index| {
479 if let Some(name) = &index.name {
480 seen.insert(name)
481 } else {
482 true
483 }
484 })
485 .partition::<Vec<_>, _>(|index| !index.default);
486
487 Either::Right(non_default.into_iter().chain(default))
488 }
489
490 fn index_for_url(&self, url: &IndexUrl) -> Option<&Index> {
492 self.indexes
493 .iter()
494 .find(|index| is_same_index(index.url(), url))
495 }
496
497 pub fn status_code_strategy_for(&self, url: &IndexUrl) -> IndexStatusCodeStrategy {
499 self.index_for_url(url).map_or(
500 IndexStatusCodeStrategy::Default,
501 Index::status_code_strategy,
502 )
503 }
504
505 pub fn simple_api_cache_control_for(&self, url: &IndexUrl) -> Option<http::HeaderValue> {
507 self.index_for_url(url)
508 .and_then(Index::simple_api_cache_control)
509 }
510
511 pub fn artifact_cache_control_for(&self, url: &IndexUrl) -> Option<http::HeaderValue> {
513 self.index_for_url(url)
514 .and_then(Index::artifact_cache_control)
515 }
516
517 pub fn exclude_newer_for(&self, url: &IndexUrl) -> Option<&ExcludeNewerOverride> {
519 self.index_for_url(url).and_then(Index::exclude_newer)
520 }
521}
522
523impl From<&IndexLocations> for uv_auth::Indexes {
524 fn from(index_locations: &IndexLocations) -> Self {
525 Self::from_indexes(index_locations.allowed_indexes().into_iter().map(|index| {
526 let mut url = index.url().url().clone();
527 url.set_username("").ok();
528 url.set_password(None).ok();
529 let mut root_url = index.url().root().unwrap_or_else(|| url.clone());
530 root_url.set_username("").ok();
531 root_url.set_password(None).ok();
532 uv_auth::Index {
533 url,
534 root_url,
535 auth_policy: index.authenticate,
536 }
537 }))
538 }
539}
540
541bitflags::bitflags! {
542 #[derive(Debug, Copy, Clone)]
543 struct Flags: u8 {
544 const NO_RANGE_REQUESTS = 1;
546 const UNAUTHORIZED = 1 << 2;
548 const FORBIDDEN = 1 << 1;
550 }
551}
552
553#[derive(Debug, Default, Clone)]
559pub struct IndexCapabilities(Arc<RwLock<FxHashMap<IndexUrl, Flags>>>);
560
561impl IndexCapabilities {
562 pub fn supports_range_requests(&self, index_url: &IndexUrl) -> bool {
564 !self
565 .0
566 .read()
567 .unwrap()
568 .get(index_url)
569 .is_some_and(|flags| flags.intersects(Flags::NO_RANGE_REQUESTS))
570 }
571
572 pub fn set_no_range_requests(&self, index_url: IndexUrl) {
574 self.0
575 .write()
576 .unwrap()
577 .entry(index_url)
578 .or_insert(Flags::empty())
579 .insert(Flags::NO_RANGE_REQUESTS);
580 }
581
582 pub fn unauthorized(&self, index_url: &IndexUrl) -> bool {
584 self.0
585 .read()
586 .unwrap()
587 .get(index_url)
588 .is_some_and(|flags| flags.intersects(Flags::UNAUTHORIZED))
589 }
590
591 pub(crate) fn set_unauthorized(&self, index_url: IndexUrl) {
593 self.0
594 .write()
595 .unwrap()
596 .entry(index_url)
597 .or_insert(Flags::empty())
598 .insert(Flags::UNAUTHORIZED);
599 }
600
601 pub fn forbidden(&self, index_url: &IndexUrl) -> bool {
603 self.0
604 .read()
605 .unwrap()
606 .get(index_url)
607 .is_some_and(|flags| flags.intersects(Flags::FORBIDDEN))
608 }
609
610 pub(crate) fn set_forbidden(&self, index_url: IndexUrl) {
612 self.0
613 .write()
614 .unwrap()
615 .entry(index_url)
616 .or_insert(Flags::empty())
617 .insert(Flags::FORBIDDEN);
618 }
619}
620
621#[cfg(test)]
622mod tests {
623 use super::*;
624 use crate::{IndexCacheControl, IndexFormat, IndexName};
625 use http::HeaderValue;
626
627 #[test]
628 fn test_index_url_parse_valid_paths() {
629 assert!(is_disambiguated_path("/absolute/path"));
631 assert!(is_disambiguated_path("./relative/path"));
633 assert!(is_disambiguated_path("../../relative/path"));
634 if cfg!(windows) {
635 assert!(is_disambiguated_path("C:/absolute/path"));
637 assert!(is_disambiguated_path(".\\relative\\path"));
639 assert!(is_disambiguated_path("..\\..\\relative\\path"));
640 }
641 }
642
643 #[test]
644 fn test_index_url_parse_ambiguous_paths() {
645 assert!(!is_disambiguated_path("index"));
647 assert!(!is_disambiguated_path("relative/path"));
649 }
650
651 #[test]
652 fn test_index_url_parse_with_schemes() {
653 assert!(is_disambiguated_path("file:///absolute/path"));
654 assert!(is_disambiguated_path("https://registry.com/simple/"));
655 assert!(is_disambiguated_path(
656 "git+https://github.com/example/repo.git"
657 ));
658 }
659
660 #[test]
661 fn fetch_indexes_deduplicates_raw_urls() {
662 let url = IndexUrl::from_str("https://index.example.com/simple").unwrap();
663 let mut first = Index::from(url.clone());
664 first.name = Some(IndexName::from_str("first").unwrap());
665 let mut second = Index::from(url);
666 second.name = Some(IndexName::from_str("second").unwrap());
667 second.default = true;
668 let locations = IndexLocations::new(vec![first, second], Vec::new(), false);
669
670 assert_eq!(locations.indexes().count(), 2);
671 assert_eq!(locations.fetch_indexes().count(), 1);
672 }
673
674 #[test]
675 fn test_cache_control_lookup() {
676 use std::str::FromStr;
677
678 use crate::IndexFormat;
679 use crate::index_name::IndexName;
680
681 let indexes = vec![
682 Index {
683 name: Some(IndexName::from_str("index1").unwrap()),
684 url: IndexUrl::from_str("https://index1.example.com/simple").unwrap(),
685 cache_control: Some(crate::IndexCacheControl {
686 api: Some(HeaderValue::from_static("max-age=300")),
687 files: Some(HeaderValue::from_static("max-age=1800")),
688 }),
689 explicit: false,
690 default: false,
691 origin: None,
692 format: IndexFormat::Simple,
693 publish_url: None,
694 authenticate: uv_auth::AuthPolicy::default(),
695 ignore_error_codes: None,
696 exclude_newer: None,
697 },
698 Index {
699 name: Some(IndexName::from_str("index2").unwrap()),
700 url: IndexUrl::from_str("https://index2.example.com/simple").unwrap(),
701 cache_control: None,
702 explicit: false,
703 default: false,
704 origin: None,
705 format: IndexFormat::Simple,
706 publish_url: None,
707 authenticate: uv_auth::AuthPolicy::default(),
708 ignore_error_codes: None,
709 exclude_newer: None,
710 },
711 ];
712
713 let index_locations = IndexLocations::new(indexes, Vec::new(), false);
714
715 let url1 = IndexUrl::from_str("https://index1.example.com/simple").unwrap();
716 assert_eq!(
717 index_locations.simple_api_cache_control_for(&url1),
718 Some(HeaderValue::from_static("max-age=300"))
719 );
720 assert_eq!(
721 index_locations.artifact_cache_control_for(&url1),
722 Some(HeaderValue::from_static("max-age=1800"))
723 );
724
725 let url2 = IndexUrl::from_str("https://index2.example.com/simple").unwrap();
726 assert_eq!(index_locations.simple_api_cache_control_for(&url2), None);
727 assert_eq!(index_locations.artifact_cache_control_for(&url2), None);
728
729 let url3 = IndexUrl::from_str("https://index3.example.com/simple").unwrap();
730 assert_eq!(index_locations.simple_api_cache_control_for(&url3), None);
731 assert_eq!(index_locations.artifact_cache_control_for(&url3), None);
732 }
733
734 #[test]
735 fn test_pytorch_default_cache_control() {
736 let indexes = vec![Index {
738 name: Some(IndexName::from_str("pytorch").unwrap()),
739 url: IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap(),
740 cache_control: None, explicit: false,
742 default: false,
743 origin: None,
744 format: IndexFormat::Simple,
745 publish_url: None,
746 authenticate: uv_auth::AuthPolicy::default(),
747 ignore_error_codes: None,
748 exclude_newer: None,
749 }];
750
751 let index_locations = IndexLocations::new(indexes, Vec::new(), false);
752
753 let pytorch_url = IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap();
754
755 assert_eq!(
756 index_locations.simple_api_cache_control_for(&pytorch_url),
757 None
758 );
759 assert_eq!(
760 index_locations.artifact_cache_control_for(&pytorch_url),
761 Some(HeaderValue::from_static(
762 "max-age=365000000, immutable, public",
763 ))
764 );
765 }
766
767 #[test]
768 fn test_pytorch_user_override_cache_control() {
769 let indexes = vec![Index {
771 name: Some(IndexName::from_str("pytorch").unwrap()),
772 url: IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap(),
773 cache_control: Some(IndexCacheControl {
774 api: Some(HeaderValue::from_static("no-cache")),
775 files: Some(HeaderValue::from_static("max-age=3600")),
776 }),
777 explicit: false,
778 default: false,
779 origin: None,
780 format: IndexFormat::Simple,
781 publish_url: None,
782 authenticate: uv_auth::AuthPolicy::default(),
783 ignore_error_codes: None,
784 exclude_newer: None,
785 }];
786
787 let index_locations = IndexLocations::new(indexes, Vec::new(), false);
788
789 let pytorch_url = IndexUrl::from_str("https://download.pytorch.org/whl/cu118").unwrap();
790
791 assert_eq!(
792 index_locations.simple_api_cache_control_for(&pytorch_url),
793 Some(HeaderValue::from_static("no-cache"))
794 );
795 assert_eq!(
796 index_locations.artifact_cache_control_for(&pytorch_url),
797 Some(HeaderValue::from_static("max-age=3600"))
798 );
799 }
800
801 #[test]
802 fn test_nvidia_default_cache_control() {
803 let indexes = vec![Index {
805 name: Some(IndexName::from_str("nvidia").unwrap()),
806 url: IndexUrl::from_str("https://pypi.nvidia.com").unwrap(),
807 cache_control: None, explicit: false,
809 default: false,
810 origin: None,
811 format: IndexFormat::Simple,
812 publish_url: None,
813 authenticate: uv_auth::AuthPolicy::default(),
814 ignore_error_codes: None,
815 exclude_newer: None,
816 }];
817
818 let index_locations = IndexLocations::new(indexes, Vec::new(), false);
819
820 let nvidia_url = IndexUrl::from_str("https://pypi.nvidia.com").unwrap();
821
822 assert_eq!(
823 index_locations.simple_api_cache_control_for(&nvidia_url),
824 None
825 );
826 assert_eq!(
827 index_locations.artifact_cache_control_for(&nvidia_url),
828 Some(HeaderValue::from_static(
829 "max-age=365000000, immutable, public",
830 ))
831 );
832 }
833}