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}