indoc_impl/
lib.rs

1extern crate proc_macro;
2
3use proc_macro2::{Span, TokenStream};
4use quote::quote;
5use syn::{Error, Lit, LitByteStr, LitStr, Result};
6use unindent::*;
7
8#[proc_macro]
9pub fn indoc(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
10    let input = TokenStream::from(input);
11
12    let unindented = match try_indoc(input) {
13        Ok(tokens) => tokens,
14        Err(err) => err.to_compile_error(),
15    };
16
17    proc_macro::TokenStream::from(unindented)
18}
19
20#[proc_macro]
21pub fn formatdoc(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
22    let input = TokenStream::from(input);
23
24    let format = match format_indoc(input, "format") {
25        Ok(tokens) => tokens,
26        Err(err) => err.to_compile_error(),
27    };
28
29    proc_macro::TokenStream::from(format)
30}
31
32#[proc_macro]
33pub fn printdoc(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
34    let input = TokenStream::from(input);
35
36    let format = match format_indoc(input, "print") {
37        Ok(tokens) => tokens,
38        Err(err) => err.to_compile_error(),
39    };
40
41    proc_macro::TokenStream::from(format)
42}
43
44fn try_indoc(input: TokenStream) -> Result<TokenStream> {
45    let len = input.clone().into_iter().count();
46    if len != 1 {
47        return Err(Error::new(
48            Span::call_site(),
49            format!(
50                "argument must be a single string literal, but got {} tokens",
51                len,
52            ),
53        ));
54    }
55    lit_indoc(input, true)
56}
57
58fn lit_indoc(input: TokenStream, raw: bool) -> Result<TokenStream> {
59    let lit = match syn::parse2::<Lit>(input) {
60        Ok(lit) => lit,
61        Err(err) => {
62            return Err(Error::new(
63                err.span(),
64                "argument must be a single string literal",
65            ));
66        }
67    };
68
69    let lit = match lit {
70        Lit::Str(lit) => {
71            let v = unindent(&lit.value());
72            Lit::Str(LitStr::new(&v, lit.span()))
73        }
74        Lit::ByteStr(lit) => {
75            if raw {
76                let v = unindent_bytes(&lit.value());
77                Lit::ByteStr(LitByteStr::new(&v, lit.span()))
78            } else {
79                return Err(Error::new_spanned(
80                    lit,
81                    "byte strings are not supported in formatting macros",
82                ));
83            }
84        }
85        otherwise => {
86            return Err(Error::new_spanned(
87                otherwise,
88                "argument must be a single string literal",
89            ));
90        }
91    };
92
93    Ok(quote!(#lit))
94}
95
96fn format_indoc(input: TokenStream, macro_name: &str) -> Result<TokenStream> {
97    let macro_name = syn::Ident::new(macro_name, Span::call_site());
98
99    let mut input = input.into_iter();
100
101    let first = std::iter::once(input.next().ok_or_else(|| {
102        Error::new(
103            Span::call_site(),
104            "unexpected end of macro invocation, expected format string",
105        )
106    })?)
107    .collect();
108    let fmt_str = lit_indoc(first, false)?;
109
110    let args: TokenStream = input.collect();
111
112    Ok(quote!(#macro_name!(#fmt_str #args)))
113}