1use std::path::{Path, PathBuf};
30
31use serde::{Deserialize, Serialize};
32
33use crate::{Error, Result};
34
35pub const PYPI_URL: &str = "https://pypi.org/simple/";
37pub const PYPI_API_URL: &str = "https://pypi.org/pypi";
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct RegistryConfig {
42 pub name: String,
44
45 pub url: String,
47
48 #[serde(skip_serializing_if = "Option::is_none")]
50 pub username: Option<String>,
51
52 #[serde(skip_serializing_if = "Option::is_none")]
54 pub password: Option<String>,
55
56 #[serde(skip_serializing_if = "Option::is_none")]
58 pub token: Option<String>,
59
60 #[serde(default)]
62 pub default: bool,
63
64 #[serde(default = "default_priority")]
66 pub priority: u32,
67}
68
69fn default_priority() -> u32 {
70 100
71}
72
73impl RegistryConfig {
74 pub fn new(name: impl Into<String>, url: impl Into<String>) -> Self {
76 Self {
77 name: name.into(),
78 url: url.into(),
79 username: None,
80 password: None,
81 token: None,
82 default: false,
83 priority: default_priority(),
84 }
85 }
86
87 pub fn pypi() -> Self {
89 Self {
90 name: "pypi".to_string(),
91 url: PYPI_URL.to_string(),
92 username: None,
93 password: None,
94 token: None,
95 default: true,
96 priority: 1000, }
98 }
99
100 pub fn with_basic_auth(mut self, username: String, password: String) -> Self {
102 self.username = Some(username);
103 self.password = Some(password);
104 self
105 }
106
107 pub fn with_token(mut self, token: String) -> Self {
109 self.token = Some(token);
110 self
111 }
112
113 pub fn resolve_credentials(&self) -> Result<ResolvedCredentials> {
115 let username = self.username.as_ref().map(|u| resolve_env_var(u));
116 let password = self.password.as_ref().map(|p| resolve_env_var(p));
117 let token = self.token.as_ref().map(|t| resolve_env_var(t));
118
119 Ok(ResolvedCredentials {
120 username,
121 password,
122 token,
123 })
124 }
125
126 pub fn has_auth(&self) -> bool {
128 self.username.is_some() || self.token.is_some()
129 }
130
131 pub fn api_url(&self) -> String {
133 if self.url.ends_with("/simple/") || self.url.ends_with("/simple") {
136 let base = self.url.trim_end_matches('/').trim_end_matches("simple");
137 format!("{}pypi", base)
138 } else {
139 self.url.clone()
140 }
141 }
142}
143
144#[derive(Debug, Clone)]
146pub struct ResolvedCredentials {
147 pub username: Option<String>,
148 pub password: Option<String>,
149 pub token: Option<String>,
150}
151
152impl ResolvedCredentials {
153 pub fn has_credentials(&self) -> bool {
155 self.username.is_some() || self.token.is_some()
156 }
157}
158
159#[derive(Debug, Clone, Default)]
161pub struct RegistryManager {
162 registries: Vec<RegistryConfig>,
164}
165
166impl RegistryManager {
167 pub fn new() -> Self {
169 Self {
170 registries: vec![RegistryConfig::pypi()],
171 }
172 }
173
174 pub fn from_configs(mut configs: Vec<RegistryConfig>) -> Self {
176 configs.sort_by_key(|r| r.priority);
178
179 if !configs.iter().any(|r| r.name == "pypi") {
181 configs.push(RegistryConfig::pypi());
182 }
183
184 Self {
185 registries: configs,
186 }
187 }
188
189 pub fn load(project_dir: &Path) -> Result<Self> {
191 let mut configs = Vec::new();
192
193 if let Some(global_configs) = load_global_config()? {
195 configs.extend(global_configs);
196 }
197
198 if let Some(project_configs) = load_project_config(project_dir)? {
200 for config in project_configs {
202 configs.retain(|c| c.name != config.name);
204 configs.push(config);
205 }
206 }
207
208 Ok(Self::from_configs(configs))
209 }
210
211 pub fn add(&mut self, config: RegistryConfig) {
213 self.registries.retain(|r| r.name != config.name);
215 self.registries.push(config);
216 self.registries.sort_by_key(|r| r.priority);
217 }
218
219 pub fn registries(&self) -> &[RegistryConfig] {
221 &self.registries
222 }
223
224 pub fn get(&self, name: &str) -> Option<&RegistryConfig> {
226 self.registries.iter().find(|r| r.name == name)
227 }
228
229 pub fn default_registry(&self) -> Option<&RegistryConfig> {
231 self.registries.iter().find(|r| r.default)
232 }
233
234 pub fn primary(&self) -> Option<&RegistryConfig> {
236 self.registries.first()
237 }
238}
239
240fn load_global_config() -> Result<Option<Vec<RegistryConfig>>> {
242 let config_path = dirs::home_dir()
243 .map(|h| h.join(".rx").join("config.toml"))
244 .unwrap_or_else(|| PathBuf::from(".rx/config.toml"));
245
246 if !config_path.exists() {
247 return Ok(None);
248 }
249
250 let content = std::fs::read_to_string(&config_path).map_err(Error::Io)?;
251
252 #[derive(Deserialize)]
253 struct GlobalConfig {
254 #[serde(default)]
255 registries: Vec<RegistryConfig>,
256 }
257
258 let config: GlobalConfig = toml::from_str(&content).map_err(Error::TomlParse)?;
259 Ok(Some(config.registries))
260}
261
262fn load_project_config(project_dir: &Path) -> Result<Option<Vec<RegistryConfig>>> {
264 let pyproject_path = project_dir.join("pyproject.toml");
265 if !pyproject_path.exists() {
266 return Ok(None);
267 }
268
269 let content = std::fs::read_to_string(&pyproject_path).map_err(Error::Io)?;
270 let doc: toml::Value = toml::from_str(&content).map_err(Error::TomlParse)?;
271
272 let registries = doc
273 .get("tool")
274 .and_then(|t| t.get("rx"))
275 .and_then(|r| r.get("registries"))
276 .and_then(|r| r.as_array());
277
278 match registries {
279 Some(arr) => {
280 let configs: Vec<RegistryConfig> = arr
281 .iter()
282 .filter_map(|v| {
283 let s = toml::to_string(v).ok()?;
284 toml::from_str(&s).ok()
285 })
286 .collect();
287 Ok(Some(configs))
288 }
289 None => Ok(None),
290 }
291}
292
293fn resolve_env_var(value: &str) -> String {
296 let mut result = value.to_string();
297
298 while let Some(start) = result.find("${") {
300 if let Some(end) = result[start..].find('}') {
301 let var_name = &result[start + 2..start + end];
302 let replacement = std::env::var(var_name).unwrap_or_default();
303 result = format!(
304 "{}{}{}",
305 &result[..start],
306 replacement,
307 &result[start + end + 1..]
308 );
309 } else {
310 break;
311 }
312 }
313
314 result
315}
316
317#[cfg(test)]
318mod tests {
319 use super::*;
320
321 #[test]
322 fn test_registry_config() {
323 let config = RegistryConfig::new("private", "https://private.pypi.org/simple/")
324 .with_basic_auth("user".to_string(), "pass".to_string());
325
326 assert_eq!(config.name, "private");
327 assert!(config.has_auth());
328 assert_eq!(config.api_url(), "https://private.pypi.org/pypi");
329 }
330
331 #[test]
332 fn test_resolve_env_var() {
333 std::env::set_var("TEST_VAR", "test_value");
334 assert_eq!(resolve_env_var("${TEST_VAR}"), "test_value");
335 assert_eq!(
336 resolve_env_var("prefix_${TEST_VAR}_suffix"),
337 "prefix_test_value_suffix"
338 );
339 std::env::remove_var("TEST_VAR");
340 }
341
342 #[test]
343 fn test_registry_manager() {
344 let mut manager = RegistryManager::new();
345 assert_eq!(manager.registries().len(), 1);
346 assert_eq!(manager.primary().unwrap().name, "pypi");
347
348 manager.add(RegistryConfig::new(
349 "private",
350 "https://private.pypi.org/simple/",
351 ));
352 assert!(manager.get("private").is_some());
353 }
354}