Skip to main content

pick/
lib.rs

1pub mod cli;
2pub mod detector;
3pub mod error;
4pub mod formats;
5pub mod output;
6pub mod selector;
7
8use cli::{Cli, InputFormat};
9use error::PickError;
10use selector::{Selector, extract};
11use serde_json::Value;
12
13pub fn run(cli: &Cli, input: &str) -> Result<String, PickError> {
14    if input.trim().is_empty() {
15        return Err(PickError::NoInput);
16    }
17
18    let selector_str = cli.selector.as_deref().unwrap_or("");
19    let selector = Selector::parse(selector_str)?;
20
21    // Determine format
22    let format = match cli.input {
23        InputFormat::Auto => detector::detect_format(input),
24        ref f => f.clone(),
25    };
26
27    // Parse and extract
28    let results = match parse_and_extract(input, &format, &selector, selector_str) {
29        Ok(r) => r,
30        Err(e) => {
31            if let Some(ref default) = cli.default {
32                return Ok(default.clone());
33            }
34            return Err(e);
35        }
36    };
37
38    // Handle empty results with --default
39    if results.is_empty() {
40        if let Some(ref default) = cli.default {
41            return Ok(default.clone());
42        }
43        return Err(PickError::KeyNotFound(selector_str.to_string()));
44    }
45
46    // --exists: just check, output nothing
47    if cli.exists {
48        return Ok(String::new());
49    }
50
51    // --count: output match count
52    if cli.count {
53        return Ok(results.len().to_string());
54    }
55
56    // --first: only first result
57    let results = if cli.first {
58        vec![results.into_iter().next().unwrap()]
59    } else {
60        results
61    };
62
63    Ok(output::format_output(&results, cli.json, cli.lines))
64}
65
66fn parse_and_extract(
67    input: &str,
68    format: &InputFormat,
69    selector: &Selector,
70    selector_str: &str,
71) -> Result<Vec<Value>, PickError> {
72    // Text format has a special fallback path
73    if *format == InputFormat::Text {
74        return parse_and_extract_text(input, selector, selector_str);
75    }
76
77    let value = parse_input(input, format)?;
78    extract(&value, selector)
79}
80
81fn parse_and_extract_text(
82    input: &str,
83    selector: &Selector,
84    selector_str: &str,
85) -> Result<Vec<Value>, PickError> {
86    let value = formats::text::parse(input)?;
87
88    // Try normal extraction first
89    if let Ok(results) = extract(&value, selector)
90        && !results.is_empty()
91    {
92        return Ok(results);
93    }
94
95    // Fallback: search for the full selector string in the text
96    if !selector_str.is_empty()
97        && let Some(found) = formats::text::search_text(input, selector_str)
98    {
99        return Ok(vec![found]);
100    }
101
102    Err(PickError::KeyNotFound(selector_str.to_string()))
103}
104
105fn parse_input(input: &str, format: &InputFormat) -> Result<Value, PickError> {
106    match format {
107        InputFormat::Json => formats::json::parse(input),
108        InputFormat::Yaml => formats::yaml::parse(input),
109        InputFormat::Toml => formats::toml_format::parse(input),
110        InputFormat::Env => formats::env::parse(input),
111        InputFormat::Headers => formats::headers::parse(input),
112        InputFormat::Logfmt => formats::logfmt::parse(input),
113        InputFormat::Csv => formats::csv_format::parse(input),
114        InputFormat::Text => formats::text::parse(input),
115        InputFormat::Auto => {
116            // Should not reach here; detect_format handles this
117            Err(PickError::UnknownFormat)
118        }
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    fn make_cli(selector: Option<&str>) -> Cli {
127        Cli {
128            selector: selector.map(String::from),
129            input: InputFormat::Auto,
130            file: None,
131            json: false,
132            raw: false,
133            first: false,
134            lines: false,
135            default: None,
136            quiet: false,
137            exists: false,
138            count: false,
139        }
140    }
141
142    #[test]
143    fn run_json_simple() {
144        let cli = make_cli(Some("name"));
145        let result = run(&cli, r#"{"name": "Alice"}"#).unwrap();
146        assert_eq!(result, "Alice");
147    }
148
149    #[test]
150    fn run_json_nested() {
151        let cli = make_cli(Some("user.email"));
152        let result = run(&cli, r#"{"user": {"email": "a@b.com"}}"#).unwrap();
153        assert_eq!(result, "a@b.com");
154    }
155
156    #[test]
157    fn run_json_array_index() {
158        let cli = make_cli(Some("items[0]"));
159        let result = run(&cli, r#"{"items": ["first", "second"]}"#).unwrap();
160        assert_eq!(result, "first");
161    }
162
163    #[test]
164    fn run_json_wildcard() {
165        let cli = make_cli(Some("items[*].name"));
166        let result = run(&cli, r#"{"items": [{"name": "a"}, {"name": "b"}]}"#).unwrap();
167        assert_eq!(result, "a\nb");
168    }
169
170    #[test]
171    fn run_yaml() {
172        let mut cli = make_cli(Some("name"));
173        cli.input = InputFormat::Yaml;
174        let result = run(&cli, "name: Alice\nage: 30").unwrap();
175        assert_eq!(result, "Alice");
176    }
177
178    #[test]
179    fn run_toml() {
180        let mut cli = make_cli(Some("package.name"));
181        cli.input = InputFormat::Toml;
182        let result = run(&cli, "[package]\nname = \"pick\"").unwrap();
183        assert_eq!(result, "pick");
184    }
185
186    #[test]
187    fn run_env() {
188        let mut cli = make_cli(Some("PORT"));
189        cli.input = InputFormat::Env;
190        let result = run(&cli, "DATABASE_URL=pg://localhost\nPORT=3000").unwrap();
191        assert_eq!(result, "3000");
192    }
193
194    #[test]
195    fn run_headers() {
196        let mut cli = make_cli(Some("content-type"));
197        cli.input = InputFormat::Headers;
198        let result = run(&cli, "Content-Type: application/json\nX-Request-Id: abc").unwrap();
199        assert_eq!(result, "application/json");
200    }
201
202    #[test]
203    fn run_logfmt() {
204        let mut cli = make_cli(Some("msg"));
205        cli.input = InputFormat::Logfmt;
206        let result = run(&cli, "level=info msg=hello status=200").unwrap();
207        assert_eq!(result, "hello");
208    }
209
210    #[test]
211    fn run_csv() {
212        let mut cli = make_cli(Some("[0].name"));
213        cli.input = InputFormat::Csv;
214        let result = run(&cli, "name,age\nAlice,30\nBob,25").unwrap();
215        assert_eq!(result, "Alice");
216    }
217
218    #[test]
219    fn run_no_selector_returns_whole() {
220        let cli = make_cli(None);
221        let result = run(&cli, r#"{"a": 1}"#).unwrap();
222        assert!(result.contains("\"a\""));
223    }
224
225    #[test]
226    fn run_empty_input() {
227        let cli = make_cli(Some("x"));
228        assert!(run(&cli, "").is_err());
229        assert!(run(&cli, "   ").is_err());
230    }
231
232    #[test]
233    fn run_key_not_found() {
234        let cli = make_cli(Some("missing"));
235        assert!(run(&cli, r#"{"a": 1}"#).is_err());
236    }
237
238    #[test]
239    fn run_default_on_missing() {
240        let mut cli = make_cli(Some("missing"));
241        cli.default = Some("fallback".into());
242        let result = run(&cli, r#"{"a": 1}"#).unwrap();
243        assert_eq!(result, "fallback");
244    }
245
246    #[test]
247    fn run_exists_found() {
248        let mut cli = make_cli(Some("a"));
249        cli.exists = true;
250        let result = run(&cli, r#"{"a": 1}"#).unwrap();
251        assert_eq!(result, "");
252    }
253
254    #[test]
255    fn run_exists_not_found() {
256        let mut cli = make_cli(Some("b"));
257        cli.exists = true;
258        assert!(run(&cli, r#"{"a": 1}"#).is_err());
259    }
260
261    #[test]
262    fn run_count() {
263        let mut cli = make_cli(Some("items[*]"));
264        cli.count = true;
265        let result = run(&cli, r#"{"items": [1, 2, 3]}"#).unwrap();
266        assert_eq!(result, "3");
267    }
268
269    #[test]
270    fn run_first() {
271        let mut cli = make_cli(Some("items[*]"));
272        cli.first = true;
273        let result = run(&cli, r#"{"items": [1, 2, 3]}"#).unwrap();
274        assert_eq!(result, "1");
275    }
276
277    #[test]
278    fn run_json_output() {
279        let mut cli = make_cli(Some("name"));
280        cli.json = true;
281        let result = run(&cli, r#"{"name": "Alice"}"#).unwrap();
282        assert_eq!(result, "\"Alice\"");
283    }
284
285    #[test]
286    fn run_lines_output() {
287        let mut cli = make_cli(Some("items"));
288        cli.lines = true;
289        let result = run(&cli, r#"{"items": ["a", "b", "c"]}"#).unwrap();
290        assert_eq!(result, "a\nb\nc");
291    }
292
293    #[test]
294    fn run_text_kv_fallback() {
295        let mut cli = make_cli(Some("name"));
296        cli.input = InputFormat::Text;
297        let result = run(&cli, "name=Alice\nage=30").unwrap();
298        assert_eq!(result, "Alice");
299    }
300
301    #[test]
302    fn run_text_search_fallback() {
303        let mut cli = make_cli(Some("error"));
304        cli.input = InputFormat::Text;
305        let result = run(&cli, "info: all good\nerror: something failed").unwrap();
306        assert_eq!(result, "something failed");
307    }
308
309    #[test]
310    fn run_auto_detect_json() {
311        let cli = make_cli(Some("x"));
312        let result = run(&cli, r#"{"x": 42}"#).unwrap();
313        assert_eq!(result, "42");
314    }
315
316    #[test]
317    fn run_auto_detect_env() {
318        let cli = make_cli(Some("PORT"));
319        let result = run(&cli, "PORT=3000\nHOST=localhost").unwrap();
320        assert_eq!(result, "3000");
321    }
322
323    #[test]
324    fn run_negative_index() {
325        let cli = make_cli(Some("[*][-1]"));
326        let result = run(&cli, "[[1,2],[3,4]]").unwrap();
327        assert_eq!(result, "2\n4");
328    }
329}