1use std::path::Path;
2use std::str::FromStr;
3
4use serde::{Deserialize, Serialize};
5use thiserror::Error;
6use url::Url;
7
8use uv_auth::{AuthPolicy, Credentials};
9use uv_redacted::DisplaySafeUrl;
10use uv_small_str::SmallString;
11
12use crate::index_name::{IndexName, IndexNameError};
13use crate::origin::Origin;
14use crate::{IndexStatusCodeStrategy, IndexUrl, IndexUrlError, SerializableStatusCode};
15
16#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Default)]
18#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
19#[serde(rename_all = "kebab-case")]
20pub struct IndexCacheControl {
21 pub api: Option<SmallString>,
23 pub files: Option<SmallString>,
25}
26
27impl IndexCacheControl {
28 pub fn simple_api_cache_control(_url: &Url) -> Option<&'static str> {
30 None
31 }
32
33 pub fn artifact_cache_control(url: &Url) -> Option<&'static str> {
35 let dominated_by_pytorch_or_nvidia = url.host_str().is_some_and(|host| {
36 host.eq_ignore_ascii_case("download.pytorch.org")
37 || host.eq_ignore_ascii_case("pypi.nvidia.com")
38 });
39 if dominated_by_pytorch_or_nvidia {
40 Some("max-age=365000000, immutable, public")
48 } else {
49 None
50 }
51 }
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize)]
55#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
56#[serde(rename_all = "kebab-case")]
57pub struct Index {
58 pub name: Option<IndexName>,
72 pub url: IndexUrl,
76 #[serde(default)]
91 pub explicit: bool,
92 #[serde(default)]
102 pub default: bool,
103 #[serde(skip)]
105 pub origin: Option<Origin>,
106 #[serde(default)]
112 pub format: IndexFormat,
113 pub publish_url: Option<DisplaySafeUrl>,
126 #[serde(default)]
135 pub authenticate: AuthPolicy,
136 #[serde(default)]
146 pub ignore_error_codes: Option<Vec<SerializableStatusCode>>,
147 #[serde(default)]
159 pub cache_control: Option<IndexCacheControl>,
160}
161
162impl PartialEq for Index {
163 fn eq(&self, other: &Self) -> bool {
164 let Self {
165 name,
166 url,
167 explicit,
168 default,
169 origin: _,
170 format,
171 publish_url,
172 authenticate,
173 ignore_error_codes,
174 cache_control,
175 } = self;
176 *url == other.url
177 && *name == other.name
178 && *explicit == other.explicit
179 && *default == other.default
180 && *format == other.format
181 && *publish_url == other.publish_url
182 && *authenticate == other.authenticate
183 && *ignore_error_codes == other.ignore_error_codes
184 && *cache_control == other.cache_control
185 }
186}
187
188impl Eq for Index {}
189
190impl PartialOrd for Index {
191 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
192 Some(self.cmp(other))
193 }
194}
195
196impl Ord for Index {
197 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
198 let Self {
199 name,
200 url,
201 explicit,
202 default,
203 origin: _,
204 format,
205 publish_url,
206 authenticate,
207 ignore_error_codes,
208 cache_control,
209 } = self;
210 url.cmp(&other.url)
211 .then_with(|| name.cmp(&other.name))
212 .then_with(|| explicit.cmp(&other.explicit))
213 .then_with(|| default.cmp(&other.default))
214 .then_with(|| format.cmp(&other.format))
215 .then_with(|| publish_url.cmp(&other.publish_url))
216 .then_with(|| authenticate.cmp(&other.authenticate))
217 .then_with(|| ignore_error_codes.cmp(&other.ignore_error_codes))
218 .then_with(|| cache_control.cmp(&other.cache_control))
219 }
220}
221
222impl std::hash::Hash for Index {
223 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
224 let Self {
225 name,
226 url,
227 explicit,
228 default,
229 origin: _,
230 format,
231 publish_url,
232 authenticate,
233 ignore_error_codes,
234 cache_control,
235 } = self;
236 url.hash(state);
237 name.hash(state);
238 explicit.hash(state);
239 default.hash(state);
240 format.hash(state);
241 publish_url.hash(state);
242 authenticate.hash(state);
243 ignore_error_codes.hash(state);
244 cache_control.hash(state);
245 }
246}
247
248#[derive(
249 Default,
250 Debug,
251 Copy,
252 Clone,
253 Hash,
254 Eq,
255 PartialEq,
256 Ord,
257 PartialOrd,
258 serde::Serialize,
259 serde::Deserialize,
260)]
261#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
262#[serde(rename_all = "kebab-case")]
263pub enum IndexFormat {
264 #[default]
266 Simple,
267 Flat,
269}
270
271impl Index {
272 pub fn from_index_url(url: IndexUrl) -> Self {
274 Self {
275 url,
276 name: None,
277 explicit: false,
278 default: true,
279 origin: None,
280 format: IndexFormat::Simple,
281 publish_url: None,
282 authenticate: AuthPolicy::default(),
283 ignore_error_codes: None,
284 cache_control: None,
285 }
286 }
287
288 pub fn from_extra_index_url(url: IndexUrl) -> Self {
290 Self {
291 url,
292 name: None,
293 explicit: false,
294 default: false,
295 origin: None,
296 format: IndexFormat::Simple,
297 publish_url: None,
298 authenticate: AuthPolicy::default(),
299 ignore_error_codes: None,
300 cache_control: None,
301 }
302 }
303
304 pub fn from_find_links(url: IndexUrl) -> Self {
306 Self {
307 url,
308 name: None,
309 explicit: false,
310 default: false,
311 origin: None,
312 format: IndexFormat::Flat,
313 publish_url: None,
314 authenticate: AuthPolicy::default(),
315 ignore_error_codes: None,
316 cache_control: None,
317 }
318 }
319
320 #[must_use]
322 pub fn with_origin(mut self, origin: Origin) -> Self {
323 self.origin = Some(origin);
324 self
325 }
326
327 pub fn url(&self) -> &IndexUrl {
329 &self.url
330 }
331
332 pub fn into_url(self) -> IndexUrl {
334 self.url
335 }
336
337 pub fn raw_url(&self) -> &DisplaySafeUrl {
339 self.url.url()
340 }
341
342 pub fn root_url(&self) -> Option<DisplaySafeUrl> {
347 self.url.root()
348 }
349
350 pub fn credentials(&self) -> Option<Credentials> {
352 if let Some(name) = self.name.as_ref() {
354 if let Some(credentials) = Credentials::from_env(name.to_env_var()) {
355 return Some(credentials);
356 }
357 }
358
359 Credentials::from_url(self.url.url())
361 }
362
363 pub fn relative_to(mut self, root_dir: &Path) -> Result<Self, IndexUrlError> {
365 if let IndexUrl::Path(ref url) = self.url {
366 if let Some(given) = url.given() {
367 self.url = IndexUrl::parse(given, Some(root_dir))?;
368 }
369 }
370 Ok(self)
371 }
372
373 pub fn status_code_strategy(&self) -> IndexStatusCodeStrategy {
375 if let Some(ignore_error_codes) = &self.ignore_error_codes {
376 IndexStatusCodeStrategy::from_ignored_error_codes(ignore_error_codes)
377 } else {
378 IndexStatusCodeStrategy::from_index_url(self.url.url())
379 }
380 }
381
382 pub fn artifact_cache_control(&self) -> Option<&str> {
384 if let Some(artifact_cache_control) = self
385 .cache_control
386 .as_ref()
387 .and_then(|cache_control| cache_control.files.as_deref())
388 {
389 Some(artifact_cache_control)
390 } else {
391 IndexCacheControl::artifact_cache_control(self.url.url())
392 }
393 }
394
395 pub fn simple_api_cache_control(&self) -> Option<&str> {
397 if let Some(api_cache_control) = self
398 .cache_control
399 .as_ref()
400 .and_then(|cache_control| cache_control.api.as_deref())
401 {
402 Some(api_cache_control)
403 } else {
404 IndexCacheControl::simple_api_cache_control(self.url.url())
405 }
406 }
407}
408
409impl From<IndexUrl> for Index {
410 fn from(value: IndexUrl) -> Self {
411 Self {
412 name: None,
413 url: value,
414 explicit: false,
415 default: false,
416 origin: None,
417 format: IndexFormat::Simple,
418 publish_url: None,
419 authenticate: AuthPolicy::default(),
420 ignore_error_codes: None,
421 cache_control: None,
422 }
423 }
424}
425
426impl FromStr for Index {
427 type Err = IndexSourceError;
428
429 fn from_str(s: &str) -> Result<Self, Self::Err> {
430 if let Some((name, url)) = s.split_once('=') {
432 if !name.chars().any(|c| c == ':') {
433 let name = IndexName::from_str(name)?;
434 let url = IndexUrl::from_str(url)?;
435 return Ok(Self {
436 name: Some(name),
437 url,
438 explicit: false,
439 default: false,
440 origin: None,
441 format: IndexFormat::Simple,
442 publish_url: None,
443 authenticate: AuthPolicy::default(),
444 ignore_error_codes: None,
445 cache_control: None,
446 });
447 }
448 }
449
450 let url = IndexUrl::from_str(s)?;
452 Ok(Self {
453 name: None,
454 url,
455 explicit: false,
456 default: false,
457 origin: None,
458 format: IndexFormat::Simple,
459 publish_url: None,
460 authenticate: AuthPolicy::default(),
461 ignore_error_codes: None,
462 cache_control: None,
463 })
464 }
465}
466
467#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd)]
469pub struct IndexMetadata {
470 pub url: IndexUrl,
472 pub format: IndexFormat,
474}
475
476impl IndexMetadata {
477 pub fn as_ref(&self) -> IndexMetadataRef<'_> {
479 let Self { url, format: kind } = self;
480 IndexMetadataRef { url, format: *kind }
481 }
482
483 pub fn into_url(self) -> IndexUrl {
485 self.url
486 }
487}
488
489#[derive(Debug, Copy, Clone)]
491pub struct IndexMetadataRef<'a> {
492 pub url: &'a IndexUrl,
494 pub format: IndexFormat,
496}
497
498impl IndexMetadata {
499 pub fn url(&self) -> &IndexUrl {
501 &self.url
502 }
503}
504
505impl IndexMetadataRef<'_> {
506 pub fn url(&self) -> &IndexUrl {
508 self.url
509 }
510}
511
512impl<'a> From<&'a Index> for IndexMetadataRef<'a> {
513 fn from(value: &'a Index) -> Self {
514 Self {
515 url: &value.url,
516 format: value.format,
517 }
518 }
519}
520
521impl<'a> From<&'a IndexMetadata> for IndexMetadataRef<'a> {
522 fn from(value: &'a IndexMetadata) -> Self {
523 Self {
524 url: &value.url,
525 format: value.format,
526 }
527 }
528}
529
530impl From<IndexUrl> for IndexMetadata {
531 fn from(value: IndexUrl) -> Self {
532 Self {
533 url: value,
534 format: IndexFormat::Simple,
535 }
536 }
537}
538
539impl<'a> From<&'a IndexUrl> for IndexMetadataRef<'a> {
540 fn from(value: &'a IndexUrl) -> Self {
541 Self {
542 url: value,
543 format: IndexFormat::Simple,
544 }
545 }
546}
547
548#[derive(Error, Debug)]
550pub enum IndexSourceError {
551 #[error(transparent)]
552 Url(#[from] IndexUrlError),
553 #[error(transparent)]
554 IndexName(#[from] IndexNameError),
555 #[error("Index included a name, but the name was empty")]
556 EmptyName,
557}
558
559#[cfg(test)]
560mod tests {
561 use super::*;
562
563 #[test]
564 fn test_index_cache_control_headers() {
565 let toml_str = r#"
567 name = "test-index"
568 url = "https://test.example.com/simple"
569 cache-control = { api = "max-age=600", files = "max-age=3600" }
570 "#;
571
572 let index: Index = toml::from_str(toml_str).unwrap();
573 assert_eq!(index.name.as_ref().unwrap().as_ref(), "test-index");
574 assert!(index.cache_control.is_some());
575 let cache_control = index.cache_control.as_ref().unwrap();
576 assert_eq!(cache_control.api.as_deref(), Some("max-age=600"));
577 assert_eq!(cache_control.files.as_deref(), Some("max-age=3600"));
578 }
579
580 #[test]
581 fn test_index_without_cache_control() {
582 let toml_str = r#"
584 name = "test-index"
585 url = "https://test.example.com/simple"
586 "#;
587
588 let index: Index = toml::from_str(toml_str).unwrap();
589 assert_eq!(index.name.as_ref().unwrap().as_ref(), "test-index");
590 assert_eq!(index.cache_control, None);
591 }
592
593 #[test]
594 fn test_index_partial_cache_control() {
595 let toml_str = r#"
597 name = "test-index"
598 url = "https://test.example.com/simple"
599 cache-control = { api = "max-age=300" }
600 "#;
601
602 let index: Index = toml::from_str(toml_str).unwrap();
603 assert_eq!(index.name.as_ref().unwrap().as_ref(), "test-index");
604 assert!(index.cache_control.is_some());
605 let cache_control = index.cache_control.as_ref().unwrap();
606 assert_eq!(cache_control.api.as_deref(), Some("max-age=300"));
607 assert_eq!(cache_control.files, None);
608 }
609}