Python-style decorator in Rust

Pythonには@でデコレータを定義する文法があります

def f(...):
    ...
f = staticmethod(f)

@staticmethod
def f(...):
    ...

この糖衣構文は存外便利なのでRustでも導入してみましょう。
デコレータは関数を受け取って関数を返すのでimpl Traitを使用すれば簡単に定義できますね。

fn logging<F>(func: F) -> impl Fn(i32) -> i32
where
    F: Fn(i32) -> i32,
{
    move |i| {
        println!("Input = {}", i);
        let out = func(i);
        println!("Output = {}", out);
        out
    }
}

ここではi32をもらってi32を返す関数を対象に、入力値と出力値を標準出力に出すデコレータを定義してみました。これは簡単につかえます：

fn add2(i: i32) -> i32 {
    i + 2
}

fn main() {
    let add2 = logging(add2);
    add2(2);
}

Pythonにおける@デコレータ構文はこの代入部分logging(add2)をあらかじめしておいてくれる糖衣構文でした。これをproc-macroで作ってみましょう：

use proc_macro::TokenStream;
use syn::*;

#[proc_macro_attribute]
pub fn deco(attr: TokenStream, func: TokenStream) -> TokenStream {
    let func = func.into();
    let attr = parse_attr(attr);
    let item_fn: ItemFn = syn::parse2(func).expect("Input is not a function");
    let vis = &item_fn.vis;
    let ident = &item_fn.ident;
    let block = &item_fn.block;

    let decl: FnDecl = *item_fn.decl;
    let inputs = &decl.inputs;
    let output = &decl.output;

    // (a: i32, b: f64) -> (a, b)のように引数だけ抜き出している
    let input_values: Vec<_> = inputs
        .iter()
        .map(|arg| match arg {
            &FnArg::Captured(ref val) => &val.pat,
            _ => unreachable!(""),
        })
        .collect();

    let caller = quote!{
        #vis fn #ident(#inputs) #output {
            let f = #attr(deco_internal);
            return f(#(#input_values,) *);

            fn deco_internal(#inputs) #output #block
        }
    };
    caller.into()
}

これを用いると次のようにデコレータをかけます：

#[deco(logging)]
fn add2(i: i32) -> i32 {
    i + 2
}

fn main() {
    add2(2);
}

このように関数に対してもproc-macroでコードを生成することができます。

問題点

Pythonの場合には*args, **kwdsの構文によって自由に引数を設定できたが、それが難しい。コードの生成時には引数のリストが取得できるので、loggingをproc-macroとして実装するならば簡単だが、Rustの関数として実装する場合は不可能となる。

最後に

コードはこちら
https://github.com/termoshtt/deco