Recently, I was working on a project that uses axum, a web framework for Rust. I wanted to generate API documentation for the project, but I found that axum doesn’t generate documentation for the routes. I had to write the documentation manually, which was a tedious task. I thought it would be nice if I could automatically generate documentation for the routes in axum.

First, I find there is a crate called aide that can generate OpenAPI documentation for Rust projects. Here is an example of how to use it:

pub fn todo_routes(state: AppState) -> ApiRouter {
    ApiRouter::new()
        .api_route("/:id/complete", put_with(complete_todo, complete_todo_docs))
         .api_route(
            "/:id",
            get_with(get_todo, get_todo_docs).delete_with(delete_todo, delete_todo_docs),
        )
        .with_state(state)
}
async fn complete_todo(
    State(app): State<AppState>,
    Path(todo): Path<SelectTodo>,
) -> impl IntoApiResponse {
    if let Some(todo) = app.todos.lock().unwrap().get_mut(&todo.id) {
        todo.complete = true;
        StatusCode::NO_CONTENT
    } else {
        StatusCode::NOT_FOUND
    }
}

fn complete_todo_docs(op: TransformOperation) -> TransformOperation {
    op.description("Complete a Todo.").response::<204, ()>()
}

async fn get_todo(
    State(app): State<AppState>,
    Path(todo): Path<SelectTodo>,
) -> impl IntoApiResponse {
    if let Some(todo) = app.todos.lock().unwrap().get(&todo.id) {
        Json(todo.clone()).into_response()
    } else {
        StatusCode::NOT_FOUND.into_response()
    }
}

fn get_todo_docs(op: TransformOperation) -> TransformOperation {
    op.description("Get a single Todo item.")
        .response_with::<200, Json<TodoItem>, _>(|res| {
            res.example(TodoItem {
                complete: false,
                description: "fix bugs".into(),
                id: Uuid::nil(),
            })
        })
        .response_with::<404, (), _>(|res| res.description("todo was not found"))
}

It need to write the documentation and specify the response code and response type manually. I think it is not a good way to write the documentation. And the api_route is not the same as axum’s route. It requires the handler to return a impl IntoApiResponse object. I have many handlers that return Result or Json object. I need to convert the Result object to impl IntoApiResponse object. It is not convenient.

At first, I decide to write a macro_rules to generate the documentation for the routes.

For a exisiting project, I need to create a doc function for each route. And I need to pass the return type of the handler and the doc description to the doc function. With the macro_rules, maybe I can get the return type of the handler manually. And I can pass the doc description to the doc function automatically.

Furthermore, I need to change the return type of the handler to impl IntoApiResponse object. All the handler will return object with into_response method. I can use the macro_rules to generate the wrapper function for the handler. The wrapper function will call the handler and convert the return value to impl IntoApiResponse object.

Here is the code:

#[macro_export]
macro_rules! impl_docs {
    ($method: tt, $des: expr, $return_v: tt, $f_name:tt, $($v:tt: $t:ty),+) => {
        paste::item! {
            #[allow(unused_variables,dead_code)]
            async fn [< $f_name _doc_wrapper>]($($v: $t),*) -> impl aide::axum::IntoApiResponse {
                use axum::response::IntoResponse;
                $f_name($($v),*).await.into_response()
            }
            #[allow(unused_variables,dead_code)]
            pub(crate) fn [< $f_name _with_doc>]() -> aide::axum::routing::ApiMethodRouter<Arc<AppState>>{
                aide::axum::routing::[< $method _with >]
                ([< $f_name _doc_wrapper>], |op: aide::transform::TransformOperation| {
                op.description($des)
                    .response::<200, frap_protocol::Json<$return_v>>()
                })
            }
        }
    };
}

Then I can use the impl_docs macro to generate the documentation for the routes. Here is an example:

impl_docs!(
    get,
    "get file info",
    FileInfoGetResponse,
    file_info_get,
    state: State<Arc<AppState>>,
    file_id: AxumPath<FileId>
);
async fn file_info_get(
    State(state): State<Arc<AppState>>,
    AxumPath(file_id): AxumPath<FileId>,
) -> AppResult<FileInfoGetResponse> {
    let res = ....;
    Ok(FileInfoGetResponse::success(ret))
}
pub(crate) fn file_router() -> ApiRouter<Arc<AppState>> {
    ApiRouter::new()
        .api_route("/:file_id/info", file_info_get_with_doc())
}

As you can see, I only need to copy the args and the return type of handler and use the impl_docs macro to generate the documentation for the routes. It is more convenient than the previous way. And I don’t need to change the return type of the handler. The impl_docs macro will generate a wrapper function for the handler. The wrapper function will call the handler and convert the return value to impl IntoApiResponse object.

But I think it is not enough. I want to get the args of the handler and the return type of the handler automatically. It’s better to use the proc_macro_attribute to generate the documentation for the routes. So I decide to write a proc_macro_attribute to generate the documentation for the routes.

In order to get the args of the handler and the return type of the handler, I need to parse the function signature. I find the syn crate can help me to parse the function signature. Here is an example of how to use it:

#[derive(Debug)]
pub(crate) struct ParamArgs {
    args: Vec<syn::MetaNameValue>,
}

impl syn::parse::Parse for ParamArgs {
    fn parse(input: syn::parse::ParseStream) -> syn::parse::Result<Self> {
        let punctuated: syn::punctuated::Punctuated<syn::MetaNameValue, syn::Token![,]> =
            <syn::punctuated::Punctuated<_, syn::Token![,]>>::parse_terminated(input)?;

        let mut args = Vec::new();
        for nv in punctuated {
            args.push(nv);
        }
        Ok(Self { args })
    }
}


#[proc_macro_attribute]
pub fn impl_docs(attr: TokenStream, mut funcs: TokenStream) -> TokenStream {
    let attr = parse_macro_input!(attr as ParamArgs);
    let item: syn::ItemFn = syn::parse(funcs.clone()).expect("expected function");
}

The ParamArgs struct is used to parse the args of the handler. Because I want the macro to be used like this:

#[impl_docs(method = "get", doc = "get file info", tag = "file,query")]

And the method="get" can be parsed by the MetaNameValue, and we can parse the doc and tag in the same way. Then we can get the attribute of the macro:

attr.args.into_iter().for_each(|arg| {
    let name = arg.path.get_ident().expect("not attr name").to_string();
    let value = match arg.value {
        syn::Expr::Lit(syn::ExprLit {
            lit: syn::Lit::Str(lit), // this is a string literal...
            ..
        }) => lit.value(),
        _ => panic!("attr value not string"),
    };
    println!("{}: {}", name, value);
});

Then we need to get the args and the return type of the handler. We can use the ItemFn to get the function signature:

#[derive(Debug)]
struct FuncDef {
    pub name: String,
    pub args_type: Vec<String>,
    pub return_type: String,
}

// input: ItemFn
let sig = input.sig;
let func_def = FuncDef {
    name: sig.ident.to_string(),
    args_type: sig
        .inputs
        .iter()
        .map(|arg| {
            let typed_arg = match arg {
                syn::FnArg::Receiver(_) => panic!("receiver args not supported(like Self)"),
                syn::FnArg::Typed(pat_type) => pat_type,
            };
            typed_arg.ty.to_token_stream().to_string()
        })
        .collect(),
    return_type: match sig.output {
        syn::ReturnType::Default => panic!("return type must be specified"),
        syn::ReturnType::Type(_, ty) => ty.to_token_stream().to_string(),
    },
};

After we get the args and the return type of the handler, we can generate the documentation for the routes. Code generation in macro is very simple. We can use the quote crate to generate the code or write code directly. For me, I prefer to write the code directly. Here is the code:

// parsed from the attribute
#[derive(Debug, Default)]
struct AideDocs {
    pub doc: Option<String>,
    pub method: Option<String>,
    pub tag: Option<Vec<String>>,
    pub summary: Option<String>,
}

fn impl_docs_fn(func_def: FuncDef, attrs: AideDocs) -> TokenStream {
    let mut op_trans = "op".to_string();
    if let Some(tag) = attrs.tag {
        op_trans.push_str(
            tag.iter()
                .map(|s| format!(".tag(\"{}\")", s))
                .collect::<Vec<String>>()
                .join("\n")
                .as_str(),
        );
    }

    if let Some(summary) = attrs.summary {
        op_trans.push_str(format!(".summary(\"{}\")", summary).as_str());
    }

    if let Some(doc) = attrs.doc {
        op_trans.push_str(format!(".description(\"{}\")", doc).as_str());
    }

    let mut return_type = func_def.return_type;
    return_type = return_type.replace(' ', "");
    // remove the wrapper type like Result, Json, etc.
    // Maybe it is not a good way to use the string to remove the wrapper type
    REMOVE_WRAPPER.iter().for_each(|wrapper| {
        if return_type.starts_with(wrapper) {
            return_type = return_type
                .replace(wrapper, "")
                .trim_start_matches('<')
                .trim_end_matches('>')
                .to_string();
        }
    });

    op_trans.push_str(format!(".response::<200, frap_protocol::Json<{}>>()", return_type).as_str());

    let res = format!(
        r###"
    pub(crate) fn {0}_with_doc() -> aide::axum::routing::ApiMethodRouter<Arc<AppState>>{{
        aide::axum::routing::{1}_with
        ( __{0}_doc_wrapper, |op: aide::transform::TransformOperation| {{
        {2}
        }})
    }}
    "###,
        func_def.name,
        attrs.method.expect("method must be specified"),
        op_trans
    );

    res.parse().expect("parse error")
}


fn impl_wrapper_fn(func_def: &FuncDef) -> TokenStream {
    let args_name = func_def
        .args_type
        .iter()
        .enumerate()
        .map(|(i, _)| format!("arg{}", i))
        .collect::<Vec<String>>();

    let func_def = format!(
        r###"
    async fn __{0}_doc_wrapper({1}) -> impl aide::axum::IntoApiResponse {{
        use axum::response::IntoResponse;
        {0}({2}).await.into_response()
    }}
    "###,
        func_def.name,
        args_name
            .iter()
            .zip(func_def.args_type.iter())
            .map(|(n, t)| format!("{}: {}", n, t))
            .collect::<Vec<String>>()
            .join(", "),
        args_name.join(", ")
    );

    func_def.parse().expect("parse error")
}

Finally, we return the generated code:

#[proc_macro_attribute]
pub fn impl_docs(attr: TokenStream, mut funcs: TokenStream) -> TokenStream {
    let attr = parse_macro_input!(attr as ParamArgs);
    let item: syn::ItemFn = syn::parse(funcs.clone()).expect("expected function");
    match docs::derive_docs(attr, item) {
        Ok(t) => {
            // need to merge the generated code
            funcs.extend(t);
            funcs
        }
        Err(e) => e.into_compile_error().into(),
    }
}

So finally we can use the impl_docs macro to generate the documentation for the routes. Here is an example:

#[impl_docs(method = "get", doc = "get file info", tag = "file,query")]
async fn file_info_get(
    State(state): State<Arc<AppState>>,
    AxumPath(file_id): AxumPath<FileId>,
) -> AppResult<FileInfoGetResponse> {
    let res = ....;
    Ok(FileInfoGetResponse::success(ret))
}
pub(crate) fn file_router() -> ApiRouter<Arc<AppState>> {
    ApiRouter::new()
        .api_route("/:file_id/info", file_info_get_with_doc())
}

It is more convenient than the previous way. And I don’t need to change or copy the args and return type of the handler. The impl_docs macro will generate a wrapper function for the handler. The wrapper function will call the handler and convert the return value to impl IntoApiResponse object. It’s convenient and easy to use. I think it is a good way to generate the documentation for the existing routers in axum. 🙂