1// #![deny(warnings)]
2
3use std::io::IsTerminal;
4use std::path::PathBuf;
5
6use clap::Parser;
7use palette::{FromColor, Lab, Srgb};
8
9mod find_dominant_colors;
10mod get_image_colors;
11mod printing;
12
13#[derive(Parser, Debug)]
14#[command(version, about = "Find the dominant colours in an image", long_about=None)]
15struct Cli {
16    /// Path to the image to inspect
17    path: PathBuf,
18
19    /// How many colours to find
20    #[arg(long = "max-colours", default_value_t = 5)]
21    max_colours: usize,
22
23    /// Find a single colour that will look best against this background
24    #[arg(long = "best-against-bg")]
25    background: Option<Srgb<u8>>,
26
27    /// Just print the hex values, not colour previews
28    #[arg(long = "no-palette")]
29    no_palette: bool,
30}
31
32fn main() {
33    let cli = Cli::parse();
34
35    let lab: Vec<Lab> = match get_image_colors::get_image_colors(&cli.path) {
36        Ok(lab) => lab,
37        Err(e) => {
38            eprintln!("{}", e);
39            std::process::exit(1);
40        }
41    };
42
43    assert!(lab.len() > 0);
44
45    let dominant_colors = find_dominant_colors::find_dominant_colors(&lab, cli.max_colours);
46
47    let selected_colors = match cli.background {
48        Some(bg) => find_dominant_colors::choose_best_color_for_bg(dominant_colors.clone(), &bg),
49        None => dominant_colors,
50    };
51
52    let rgb_colors = selected_colors
53        .iter()
54        .map(|c| Srgb::from_color(*c).into_format())
55        .collect::<Vec<Srgb<u8>>>();
56
57    // Should we print with colours in the terminal, or just sent text?
58    //
59    // When I created this tool, I had a `--no-palette` flag to suppress the
60    // terminal colours, but I've since realised that I can look for the
61    // presence of a TTY and disable colours if we're not in a terminal,
62    // even if the user hasn't passed `--no-palette`.
63    //
64    // I'm keeping the old flag for backwards compatibility, but I might
65    // retire it in a future v2 update.
66    //
67    // Note: because of the difficulty of simulating a TTY in automated tests,
68    // this isn't tested properly -- but I'll notice quickly if this breaks!
69    let include_bg_color = if cli.no_palette {
70        false
71    } else if std::io::stdout().is_terminal() {
72        true
73    } else {
74        false
75    };
76
77    for c in rgb_colors {
78        printing::print_color(c, &cli.background, include_bg_color);
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use assert_cmd::Command;
85    use predicates::prelude::*;
86
87    // Note: for the purposes of these tests, I mostly trust the k-means code
88    // provided by the external library.
89
90    #[test]
91    fn it_prints_the_colour() {
92        Command::cargo_bin("dominant_colours")
93            .unwrap()
94            .arg("./src/tests/red.png")
95            .assert()
96            .success()
97            .stdout("#fe0000\n")
98            .stderr("");
99    }
100
101    #[test]
102    fn it_can_look_at_png_images() {
103        assert_gets_colours_from_image("./src/tests/red.png", "#fe0000\n");
104    }
105
106    #[test]
107    fn it_can_look_at_jpeg_images() {
108        assert_gets_colours_from_image("./src/tests/black.jpg", "#000000\n");
109    }
110
111    #[test]
112    fn it_can_look_at_static_gif_images() {
113        assert_gets_colours_from_image("./src/tests/yellow.gif", "#fffb00\n");
114    }
115
116    #[test]
117    fn it_can_look_at_tiff_images() {
118        assert_gets_colours_from_image("./src/tests/green.tiff", "#04ff02\n");
119    }
120
121    #[test]
122    fn it_omits_the_escape_codes_with_no_palette() {
123        assert_gets_colours_from_image("./src/tests/red.png", "#fe0000\n");
124    }
125
126    #[test]
127    fn it_defaults_to_five_colours() {
128        let has_five_lines = predicate::str::is_match(r"^(#[a-f0-9]{6}\n){5}$").unwrap();
129
130        Command::cargo_bin("dominant_colours")
131            .unwrap()
132            .arg("./src/tests/noise.jpg")
133            .assert()
134            .success()
135            .stdout(has_five_lines)
136            .stderr("");
137    }
138
139    #[test]
140    fn it_lets_you_choose_the_max_colours() {
141        let has_eight_lines = predicate::str::is_match(r"^(#[a-f0-9]{6}\n){8}$").unwrap();
142
143        Command::cargo_bin("dominant_colours")
144            .unwrap()
145            .args(&["./src/tests/noise.jpg", "--max-colours=8"])
146            .assert()
147            .success()
148            .stdout(has_eight_lines)
149            .stderr("");
150    }
151
152    // The image created in the next two tests was created with the
153    // following command:
154    //
155    //      convert \
156    //        -delay 200 \
157    //        -loop 10 \
158    //        -dispose previous \
159    //        red.png blue.png \
160    //        red.png blue.png \
161    //        red.png blue.png \
162    //        red.png blue.png \
163    //        animated_squares.gif
164    //
165    // It creates an animated GIF that has alternating red/blue frames.
166
167    #[test]
168    fn it_looks_at_multiple_frames_in_an_animated_gif() {
169        assert_gets_colours_from_image("./src/tests/animated_squares.gif", "#0200ff\n#ff0000\n");
170    }
171
172    #[test]
173    fn it_looks_at_multiple_frames_in_an_animated_gif_uppercase() {
174        assert_gets_colours_from_image(
175            "./src/tests/animated_upper_squares.GIF",
176            "#0200ff\n#ff0000\n",
177        );
178    }
179
180    // This is an animated WebP that has alternating red/blue frames.
181    //
182    // It needs to look at multiple frames to see both colours.
183    #[test]
184    fn it_looks_at_multiple_frames_in_an_animated_webp() {
185        assert_gets_colours_from_image(
186            "./src/tests/animated_squares.webp",
187            "#0200ff\n#ff0100\n#ff0002\n",
188        );
189    }
190
191    /// Get the dominant colours for an image, and check it succeeds
192    /// with the given stdout.
193    fn assert_gets_colours_from_image(
194        path: &str,
195        expected_stdout: &str,
196    ) -> assert_cmd::assert::Assert {
197        Command::cargo_bin("dominant_colours")
198            .unwrap()
199            .arg(path)
200            .assert()
201            .success()
202            .stdout(predicate::eq(expected_stdout))
203            .stderr("")
204    }
205
206    #[test]
207    fn it_fails_if_you_pass_an_invalid_max_colours() {
208        Command::cargo_bin("dominant_colours")
209            .unwrap()
210            .args(&["./src/tests/red.png", "--max-colours=NaN"])
211            .assert()
212            .failure()
213            .code(2)
214            .stdout("")
215            .stderr("error: invalid value 'NaN' for '--max-colours <MAX_COLOURS>': invalid digit found in string\n\nFor more information, try '--help'.\n");
216    }
217
218    #[test]
219    fn it_fails_if_you_pass_an_nonexistent_file() {
220        assert_file_fails_with_error(
221            "./doesnotexist.jpg",
222            "No such file or directory (os error 2)\n",
223        );
224    }
225
226    #[test]
227    fn it_fails_if_you_pass_an_nonexistent_gif() {
228        assert_file_fails_with_error(
229            "./doesnotexist.gif",
230            "No such file or directory (os error 2)\n",
231        );
232    }
233
234    #[test]
235    fn it_fails_if_you_pass_a_non_image_file() {
236        assert_file_fails_with_error(
237            "./README.md",
238            "Unable to determine image format from file extension\n",
239        );
240    }
241
242    #[test]
243    fn it_fails_if_you_pass_an_unsupported_image_format() {
244        assert_file_fails_with_error(
245            "./src/tests/orange.heic",
246            "Unable to determine image format from file extension\n",
247        );
248    }
249
250    #[test]
251    fn it_fails_if_you_pass_a_malformed_image() {
252        assert_file_fails_with_error(
253            "./src/tests/malformed.txt.png",
254            "Format error decoding Png: Invalid PNG signature.\n",
255        );
256    }
257
258    #[test]
259    fn it_fails_if_you_pass_a_malformed_gif() {
260        assert_file_fails_with_error(
261            "./src/tests/malformed.txt.gif",
262            "Format error decoding Gif: malformed GIF header\n",
263        );
264    }
265
266    #[test]
267    fn it_fails_if_you_pass_a_malformed_webp() {
268        assert_file_fails_with_error(
269            "./src/tests/malformed.txt.webp",
270            "Format error decoding WebP: Invalid Chunk header: [52, 49, 46, 46]\n",
271        );
272    }
273
274    #[test]
275    fn it_fails_if_you_pass_a_path_without_a_file_extension() {
276        assert_file_fails_with_error(
277            "./src/tests/noextension",
278            "Path has no file extension, so could not determine image format\n",
279        );
280    }
281
282    /// Try to get the dominant colours for a file, and check it fails
283    /// with the given error message.
284    fn assert_file_fails_with_error(
285        path: &str,
286        expected_stderr: &str,
287    ) -> assert_cmd::assert::Assert {
288        Command::cargo_bin("dominant_colours")
289            .unwrap()
290            .arg(path)
291            .assert()
292            .failure()
293            .code(1)
294            .stdout("")
295            .stderr(predicate::eq(expected_stderr))
296    }
297
298    #[test]
299    fn it_chooses_the_right_color_for_a_dark_background() {
300        Command::cargo_bin("dominant_colours")
301            .unwrap()
302            .args(&[
303                "src/tests/stripes.png",
304                "--max-colours=5",
305                "--best-against-bg=#222",
306            ])
307            .assert()
308            .success()
309            .stdout("#d4fb79\n")
310            .stderr("");
311    }
312
313    #[test]
314    fn it_chooses_the_right_color_for_a_light_background() {
315        Command::cargo_bin("dominant_colours")
316            .unwrap()
317            .args(&[
318                "src/tests/stripes.png",
319                "--max-colours=5",
320                "--best-against-bg=#fff",
321            ])
322            .assert()
323            .success()
324            .stdout("#693900\n")
325            .stderr("");
326    }
327
328    #[test]
329    fn it_prints_the_version() {
330        // Match strings like `dominant_colours 1.2.3`
331        let is_version_string =
332            predicate::str::is_match(r"^dominant_colours [0-9]+\.[0-9]+\.[0-9]+\n$").unwrap();
333
334        Command::cargo_bin("dominant_colours")
335            .unwrap()
336            .arg("--version")
337            .assert()
338            .success()
339            .stdout(is_version_string)
340            .stderr("");
341    }
342}