Skip to main content

dspk_hub/catalog/
mod.rs

1//! Catalogs of published artifacts, as data: what each file is, where it is downloaded from, how it
2//! is verified (SHA-256, size) and whom to credit.
3//!
4//! One JSON document per domain, compiled in by its feature (`synapse`: [`synapse`]).
5//! [`Catalog::load_default`] merges the compiled-in documents with an optional user file named by
6//! [`CATALOG_ENV`] (entries with the same id replace the built-in ones, others are added).
7
8use std::path::Path;
9
10use anyhow::{Context, Result};
11use serde::{Deserialize, Serialize};
12
13use crate::Artifact;
14
15#[cfg(feature = "synapse")]
16pub mod synapse;
17
18/// Environment variable naming a JSON catalog merged over the built-in ones.
19pub const CATALOG_ENV: &str = "DSPK_HUB_CATALOG";
20
21/// Version of the catalog documents this crate reads.
22pub const CATALOG_SCHEMA_VERSION: u32 = 2;
23
24/// A paper to cite.
25#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
26pub struct CreditPaper {
27    pub title: String,
28    pub authors: Vec<String>,
29    pub venue: String,
30    pub year: u32,
31    pub doi: String,
32}
33
34/// The code the artifact comes from.
35#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
36pub struct CreditCode {
37    pub url: String,
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub license: Option<String>,
40    #[serde(default)]
41    pub version: String,
42}
43
44/// Whom to credit for an artifact.
45#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
46pub struct Credit {
47    #[serde(default, skip_serializing_if = "Option::is_none")]
48    pub paper: Option<CreditPaper>,
49    pub code: CreditCode,
50    /// Licence of the file itself, when its publisher states one (else `None`: never re-hosted).
51    #[serde(default, skip_serializing_if = "Option::is_none")]
52    pub file_license: Option<String>,
53    #[serde(default)]
54    pub notes: String,
55}
56
57impl Credit {
58    /// One-line citation: authors (year), title, venue, DOI, then the code.
59    pub fn citation(&self) -> String {
60        let code = match &self.code.license {
61            Some(l) => format!("Code: {} ({}, {l})", self.code.url, self.code.version),
62            None => format!("Code: {} ({})", self.code.url, self.code.version),
63        };
64        match &self.paper {
65            Some(p) => {
66                let authors = match p.authors.as_slice() {
67                    [] => String::new(),
68                    [one] => one.clone(),
69                    [first, second] => format!("{first} & {second}"),
70                    [first, .., last] => format!("{first} … {last}"),
71                };
72                format!("{authors} ({}). {}. {}. https://doi.org/{}. {code}", p.year, p.title, p.venue, p.doi)
73            }
74            None => code,
75        }
76    }
77}
78
79/// One published artifact.
80#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
81pub struct CatalogEntry {
82    /// Unique id `"<family>/<slug>"` (e.g. `"kilosort4/wtemp-v1"`).
83    pub id: String,
84    /// Human-readable title.
85    pub name: String,
86    /// Family (sorter or model) the artifact belongs to.
87    pub family: String,
88    /// What the file is (`"templates"`, `"weights"`, …), shown in listings.
89    pub kind: String,
90    #[serde(default)]
91    pub description: String,
92    /// File name in the cache (the publisher's name, e.g. `wTEMP.npz`).
93    pub file_name: String,
94    /// Download URI (`hf://`, `gh://`, `zenodo://`, `https://`, `file://`).
95    pub url: String,
96    /// Lowercase hex SHA-256 (required: an entry without one cannot be pulled).
97    pub sha256: String,
98    pub size_bytes: u64,
99    pub credit: Credit,
100}
101
102impl CatalogEntry {
103    /// What [`crate::Hub`] downloads and verifies for this entry.
104    pub fn artifact(&self) -> Artifact {
105        Artifact {
106            id: self.id.clone(),
107            family: self.family.clone(),
108            file_name: self.file_name.clone(),
109            url: self.url.clone(),
110            sha256: self.sha256.clone(),
111            size_bytes: self.size_bytes,
112        }
113    }
114}
115
116/// A catalog document: `{"schema_version": 2, "artifacts": [...]}`.
117#[derive(Debug, Clone, Serialize, Deserialize)]
118struct CatalogDocument {
119    schema_version: u32,
120    artifacts: Vec<CatalogEntry>,
121}
122
123/// Entries of one or more catalog documents.
124#[derive(Debug, Clone, Default)]
125pub struct Catalog {
126    entries: Vec<CatalogEntry>,
127}
128
129impl Catalog {
130    /// The compiled-in catalogs (one per enabled feature), plus the file named by [`CATALOG_ENV`].
131    pub fn load_default() -> Result<Self> {
132        let mut catalog = Self::default();
133        #[cfg(feature = "synapse")]
134        catalog.merge(Self::from_json_str(synapse::SYNAPSE_CATALOG_JSON).context("built-in synapse catalog")?);
135        if let Some(path) = std::env::var_os(CATALOG_ENV).filter(|p| !p.is_empty()) {
136            catalog.merge_file(Path::new(&path))?;
137        }
138        Ok(catalog)
139    }
140
141    /// Parses one catalog document.
142    pub fn from_json_str(json: &str) -> Result<Self> {
143        let doc: CatalogDocument = serde_json::from_str(json).context("invalid catalog JSON")?;
144        anyhow::ensure!(
145            doc.schema_version == CATALOG_SCHEMA_VERSION,
146            "catalog schema_version {} is not supported (this dspk-hub reads {CATALOG_SCHEMA_VERSION})",
147            doc.schema_version
148        );
149        Ok(Self { entries: doc.artifacts })
150    }
151
152    /// Merges the catalog file at `path` (its entries replace those with the same id).
153    pub fn merge_file(&mut self, path: &Path) -> Result<()> {
154        let text = std::fs::read_to_string(path).with_context(|| format!("catalog {}", path.display()))?;
155        self.merge(Self::from_json_str(&text).with_context(|| format!("catalog {}", path.display()))?);
156        Ok(())
157    }
158
159    /// Merges `other`: its entries replace those with the same id, the others are appended.
160    pub fn merge(&mut self, other: Catalog) {
161        for incoming in other.entries {
162            match self.entries.iter_mut().find(|e| e.id == incoming.id) {
163                Some(existing) => *existing = incoming,
164                None => self.entries.push(incoming),
165            }
166        }
167    }
168
169    pub fn entries(&self) -> &[CatalogEntry] {
170        &self.entries
171    }
172
173    pub fn get(&self, id: &str) -> Option<&CatalogEntry> {
174        self.entries.iter().find(|e| e.id == id)
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181
182    #[test]
183    fn merge_replaces_by_id_and_appends_new_ids() {
184        let doc = |id: &str, url: &str| {
185            format!(
186                r#"{{"schema_version": 2, "artifacts": [{{"id": "{id}", "name": "n", "family": "f", "kind": "k", "file_name": "a.bin",
187                "url": "{url}", "sha256": "{}", "size_bytes": 1, "credit": {{"code": {{"url": "https://example.org"}}}}}}]}}"#,
188                "0".repeat(64)
189            )
190        };
191        let mut catalog = Catalog::from_json_str(&doc("f/a", "https://one")).unwrap();
192        catalog.merge(Catalog::from_json_str(&doc("f/a", "https://two")).unwrap());
193        catalog.merge(Catalog::from_json_str(&doc("f/b", "https://three")).unwrap());
194        assert_eq!(catalog.entries().len(), 2);
195        assert_eq!(catalog.get("f/a").unwrap().url, "https://two");
196        assert!(Catalog::from_json_str(r#"{"schema_version": 1, "artifacts": []}"#).is_err(), "other schema versions are refused");
197    }
198}