auto_import/
lib.rs

1use proc_macro::TokenStream;
2use std::collections::BTreeSet;
3use std::io::{stderr, stdout, Write};
4use std::process::{exit, Command};
5use std::str::FromStr;
6use std::sync::atomic::{AtomicBool, Ordering::Relaxed};
7
8#[proc_macro]
9pub fn magic(input: TokenStream) -> TokenStream {
10    assert!(
11        input.is_empty(),
12        "auto_import::magic!() takes no arguments!"
13    );
14
15    static ONCE: AtomicBool = AtomicBool::new(false);
16
17    if ONCE.swap(true, Relaxed) {
18        panic!("don't call auto_import::magic!() more than once per crate!");
19    }
20
21    if let Ok(x) = std::env::var("autoimport") {
22        return TokenStream::from_str(&x).unwrap();
23    }
24
25    let mut imports = BTreeSet::new();
26
27    let mut attempts = 0;
28    loop {
29        attempts += 1;
30        let mut change = false;
31        let mut args = std::env::args_os();
32        let out = Command::new(args.next().unwrap())
33            .args(args.filter(|arg| {
34                arg.to_str()
35                    .map_or(true, |s| !s.starts_with("--error-format="))
36            }))
37            .arg("--error-format=json")
38            .env(
39                "autoimport",
40                imports.iter().map(String::as_str).collect::<String>(),
41            )
42            .output()
43            .unwrap();
44        if out.status.success() {
45            exit(0);
46        }
47        for line in std::str::from_utf8(&out.stderr)
48            .unwrap()
49            .lines()
50            .filter(|l| l.starts_with('{'))
51        {
52            if let Ok(d) = json::parse(line) {
53                for c in d["children"].members() {
54                    let suggestion = c["spans"][0]["suggested_replacement"]
55                        .as_str()
56                        .unwrap_or_default();
57                    if c["spans"][0]["text"].is_empty()
58                        && suggestion.starts_with("use ")
59                        && imports.insert(suggestion.to_string())
60                    {
61                        println!("\x1b[1;32m   Injecting\x1b[m {}", suggestion.trim());
62                        change = true;
63                    }
64                }
65            }
66        }
67        if !change || attempts == 10 {
68            stderr().write_all(&out.stderr).unwrap();
69            stdout().write_all(&out.stdout).unwrap();
70            exit(out.status.code().unwrap_or(1));
71        }
72    }
73}