如何基于属性宏构造 Rust 装饰器?

2,078 阅读5分钟

什么是装饰器?

熟悉 PythonTypeScript 语言的朋友,大抵都是了解装饰器(Decorator)及其在基于切片编程范式(AOP)中的应用。当然,就算不了解也没关系,我们可以从 23 种设计模式中的装饰器模式看起。

装饰器模式

对于装饰器模式(Decorator),Refactoring 对其有清晰的定义:

Decorator is a structural design pattern that lets you attach new behaviors to objects by placing these objects inside special wrapper objects that contain the behaviors.

所谓装饰,就是给特定的对象赋予新的行为;至于如何装饰,很简单,设计一个包含新行为的对象容器,把想要装饰的对象置于其中即可。也就是说,装饰器就是一个对象容器。那么,装饰器模式具体是如何应用的呢?假定我们需要一个具备某些操作的组件,那么我们很自然的就会这么写:

interface Component {
    operation: () => string;
}

class Earphone implements Component {
    public operation(): string {
        return "Listen Music";
    }
}

就目前来说,Earphone 组件的用途只有一个,那就是听音乐;事实上,Earphone 组件还应该包括接听电话的功能。因此,我们可以尝试通过装饰器的方式赋予其新的用途。前文提到,装饰器是一个对象容器,能够将被装饰的对象置于其中,那么:

class Decorator implements Component {
    protected component: Component;

    constructor(component: Component) {
        this.component = component;
    }

    public operation(): string {
        return this.component.operation();
    }
}

基类 Decorator 能做的事也很简单,就是以容器的形式将 Component 装进来。那如何赋予装进来的对象新的行为呢?

class EarphoneDecorator extends Decorator {
    public operation(): string {
        return `Answer Phone | ${super.operation()}`
    }
}

const earphone: Earphone = new Earphone();
console.log(earphone.operation());    // 'Listen Music'

const airpods: EarphoneDecorator = new EarphoneDecorator(earphone);
console.log(airpods.operation());    // 'Answer Phone | Listen Music'

通过扩展基类 Decorator 我们得到了一个用于装饰 Earphone 组件的装饰器,以一层包一层的方式,接管 operation 方法的执行,并完成装饰。故而,Decorator 有时也被称为 Wrapper。不论是在 Python 或是 TypeScript 语言中,设计装饰器,本质上就是设计一层 Wrapper。

Python 中的装饰器

介绍装饰器模式,离不开面对对象编程范式。Python 语言中的装饰器类型有许多,最常见的便是函数装饰器。在函数式编程(FP)范式中,存在一个高阶函数的概念,即可传入函数作为参数或将函数作为输出返回的函数。比如:

def create_adder(x: int) -> Callable[[int], int]:
    def adder(y: int) -> int:
        return x + y

    return adder

adder_15 = create_adder(15)
print(adder_15(10))    # Out: 25

如上所示,函数 create_adder 返回函数 adder,符合高阶函数的定义。此外,若是传入参数是一个函数,则有:

def decorator(func: Callable[..., Any]) -> Callable[[], Any]:
    def wrapped():
        print("Before run func()")
	func()
	print("After run func()")

    return wrapped

def greet() -> None:
    print("Hello, decorator")

greet = decorator(greet)

# Out:
# Before run func()
# Hello, decorator
# After run func()
greet()

同样作为高阶函数,函数 decorator 需传入一个函数作为参数,同时返回一个函数作为输出,由此便可在函数执行时,执行前后各打印日志。本质上,装饰器就是一个以函数为输入,同时以函数为输出的高阶函数。因此,以 Pythonic 的方式重写 greet 函数,则有:

@decorator
def greet() -> None:
    print("Hello, decorator")

# Out:
# Before run func()
# Hello, decorator
# After run func()
greet()

TypeScript 中的装饰器

相较于 Python 中的装饰器,TypeScript 中的装饰器显得略微复杂一些。不过,两者在原理上依旧是类似的。如何将一个装饰器应用到一个声明上?官方文档给的建议的是写一个装饰器工厂函数,比如:

function enumerable(value: boolean) { // 装饰器工厂函数
    // 返回装饰器
    return function (target: any, propertyKey: string, descriptor: PropertyDescriptor) {
        descriptor.enumerable = value;
    };
}

不难发现,装饰器工厂函数本质上依旧是一个返回返回函数的高阶函数。不同的是,TS 中的装饰器工厂函数所返回的函数,带有特定的参数。我们可以在具体的使用场景中打印这些参数:

class Greeter {
    greeting: string;
    constructor(message: string) {
        this.greeting = message;
    }

    @enumerable(false)
    greet() {
        return "Hello, " + this.greeting;
    }
}

/**
 * target: Greet {}
 * propertyKey: greet
 * descriptor: { "writable": true, "enumerable": false, "configurable": true }
 */

其中,参数 target 表示为装饰器所修饰的目标对象,参数 propertyKey 表示为方法装饰器所修饰的属性名,参数 descriptor 则表示属性描述。借助如上三个参数,装饰器可以为所装饰的对象赋予新的行为。

那么,在 Rust 中,是否存在如 PythonTypeScript 中的装饰器呢?看起来似乎并不存在。不过我们可以基于过程宏(Procedural macros)为 Rust 实现装饰器。

过程宏(proc_marco)

借助过程宏(procedural macros),我们能够使用 Rust 语法创建句法扩展,并如函数一般执行。在 Rust 中,存在三种过程宏:

  • 函数式宏,形如函数,并像函数一般调用,比如:custom!(...)
  • 派生宏,能够为结构体(struct)、枚举(enum)等实现函数或特征(Trait);
  • 属性宏,用于结构体、字段、函数等,为其指定属性。

不难发现,属性宏与我们在 PythonTypeScript 中所见到的装饰器十分相似,我们或可以通过借助属性宏,实现 Rust 中的装饰器。如何实现一个属性宏呢?很简单:

use proc_macro::TokenStream;

#[proc_macro_attribute]
pub fn attr_macro(attr: TokenStream, item: TokenStream) -> TokenStream {
    // code ...
}

如上所示,欲实现一个属性宏,我们需要从 proc_macro 引入 TokenStream,其表示:

representing an abstract stream of tokens, or, more specifically, a sequence of token trees.

我们可以借其遍历字符树,或者将字符树转换为数据流。接下来,我们需要像往常一般设计一个函数,同时在函数上加上 #[proc_macro_attribute] 用以表明这是一个属性宏函数;其需传入两个类型同为 TokenStream 的参数,attr 表示属性宏函数的参数,item 表示被赋予属性宏的函数本身。比如:

// my-macro/src/lib.rs
extern crate proc_macro;
use proc_macro::TokenStream;

#[proc_macro_attribute]
pub fn show_streams(attr: TokenStream, item: TokenStream) -> TokenStream {
    println!("attr: \"{}\"", attr.to_string());
    println!("item: \"{}\"", item.to_string());
    item
}

// src/lib.rs
extern crate my_macro;

use my_macro::show_streams;

// Example: Basic function
#[show_streams]
fn invoke1() {}
// out: attr: ""
// out: item: "fn invoke1() { }"

// Example: Attribute with input
#[show_streams(bar)]
fn invoke2() {}
// out: attr: "bar"
// out: item: "fn invoke2() {}"

// Example: Multiple tokens in the input
#[show_streams(multiple => tokens)]
fn invoke3() {}
// out: attr: "multiple => tokens"
// out: item: "fn invoke3() {}"

// Example:
#[show_streams { delimiters }]
fn invoke4() {}
// out: attr: "delimiters"
// out: item: "fn invoke4() {}"

借助 attritem,我们便可以做很多事,包括实现 Rust 中的装饰器。

Rust 中的装饰器

我们已经知道装饰器大概是个什么东西,也知道过程宏大概能做什么。如果不使用宏,要实现一个装饰器,本质上仍然是实现一个高阶函数。比如,我们创建一个最最简单的加法函数:

pub fn add(i: i32) -> i32 {
    i + 1
}

函数 add 接受一个 32 位整数最为参数,加一返回。假如此时我们对其进行埋点,又不更改原函数内部的逻辑,那么可以将其传给一个高阶函数,并给在高阶函数内部执行埋点逻辑:

pub 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
    }
}

如此,我们通过创建一个高阶函数实现了 Rust 中的装饰。接下来,我们可以尝试使用过程宏实现装饰器。此时,我们需要借助两个库:

[dependencies]
syn = { version="1.0.103", features=["full"] }
quote = "1.0"

其中,

  • syn 库的作用是将 Rust 字符流解析为 Rust 源代码中的句法树;
  • quote 库则是提供了一个 quote! 宏,将 Rust 句法树数据结构转换为源代码;

回想属性宏函数的签名,其传入的参数类型和函数输出类型均是 TokenStream,亦可以更深刻地理解,属性宏的本质,就是将输入源代码,经过一番处理后,输出源代码。正因如此,我们能够通过过程宏实现装饰器。需要注意的是,在引入 syn 库之前,我们需要在 Cargo.toml 中为其指定 full 功能版本,如此方能使用所有库内的函数和结构体。

我们先从简单的开始,创建一个新项目,同时为了使项目结构更加清晰,我们在新项目内再创建一个库(lib):

$ cargo new rust-decorator
$ cd rust-decorator
$ cargo new macros --lib

为了能在 rust-decorator 调用 macros 内的宏,我们需要在项目根目录的 Cargo.toml 文件中加入对应的依赖:

# rust-decorator/Cargo.toml
[dependencies]
macros = { path = "./macros" }

再来看 macros 库,为了能够编写过程宏,我们需要在对应的 Cargo.toml 引入相关依赖:

# rust-decorator/macros/Cargo.toml
[lib]
proc-macro = true

[dependencies]
syn = { version = "1.0.103", features = ["full", "extra-traits"] }
quote = "1.0"

对于 syn 库,这里往 features 内加了 extra-traits ,如此在代码开发过程中,我们能够打印查看类如 syn::TokenStream / syn::Visibility 等类型的变量。

接下来,我们就可以正式开始编写过程宏啦!首先,我们暂时忽视装饰器传入的参数,仅为函数实现一个 logging 装饰器,其效果是在函数执行前后,打印日志。那么首先,先写好函数签名:

#[proc_macro_attribute]
pub fn logging(attr: TokenStream, item: TokenStream) -> TokenStream {}

结合上文对过程宏的描述,我们知道 attr 可以作为装饰器的参数,item 则代表装饰器所装饰的函数。在实现装饰器之前,我们可以先简单学习一下前文提到的两个库,如何处理一个函数。我们可以借助 syn 库解析函数,接着使用 quote 库提供的宏重新拼接函数:

extern crate proc_macro;

use proc_macro::TokenStream;
use quote::quote;
use syn::{ItemFn, parse_macro_input};

#[proc_macro_attribute]
pub fn logging(attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);
		
    let vis = &input.vis;
    let ident = &input.sig.ident;
    let block = &input.block;

    let gen = quote! {
	#vis fn #ident() {
            println!("Ah Huh~");
            #block
        }
    };

    gen.into()
}

由于我们要处理的是一个函数,因此需要将 item 解析为 ItemFn 类型,该类型内部描述了关于函数的一切,比如可见性,函数签名等。要通过 quote 宏重组一个函数,那我们需要拿到函数的可见性,函数名称及函数体,最终在重组函数时,我们将自己想要的逻辑插入其中。了解的这两个库的基本使用姿势之后,我们可以开始尝试正式实现一个装饰器啦!

按照我们学到的姿势,我们先整一个过程宏的基本写法:

#[proc_macro_attribute]
pub fn deco(attr: TokenStream, item: TokenStream) -> TokenStream {
    let decorator = parse_macro_input!(attr as Ident);
    let decoratee = parse_macro_input!(item as ItemFn);

    let caller = quote! {
	// ...TODO
    };
    caller.into()
}

我们知道,装饰器的本质就是在一个函数外再套一个新函数,新函数内部执行函数的同时,顺带执行自身的逻辑,即“装饰品”。同样以前文中的 logging 函数为例,可知参数 attr 表示一个函数名称,item 则是函数,因此可以将前者解析为一个 Ident,后者解析为 ItemFn。接着从 ItemFn 中获取函数名称,输入输出及其可见性,用于创建高阶函数:

use syn::{Ident, FnArg}

#[proc_macro_attribute]
pub fn deco(attr: TokenStream, item: TokenStream) -> TokenStream {
    let decorator = parse_macro_input!(attr as Ident);
    let decoratee = parse_macro_input!(item as ItemFn);

    let vis = &decoratee.vis;
    let ident = &decoratee.sig.ident;
    let block = &decoratee.block;
    let inputs = &decoratee.sig.inputs;
    let output = &decoratee.sig.output;

    let caller = quote! {
	#vis fn #ident(#inputs) #output {
            // ...TODO
            fn original_fn(#inputs) #output #block
	}
    };
    caller.into()
}

显然,我们只需要在高阶函数内部执行函数 original_fn 即可完成装饰起的封装:

#[proc_macro_attribute]
pub fn deco(attr: TokenStream, item: TokenStream) -> TokenStream {
    let decorator = parse_macro_input!(attr as Ident);
    // ...

    let caller = quote! {
        #vis fn #ident(#inputs) #output {
            let func = #decorator(original_fn);
            return func(#inputs);

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

如果直接运行上述代码,那么会遇到报错:

error[E0658]: type ascription is experimental.
  --> src/main.rs:16:8
   |
 16| fn add(i: i32) -> i32 {
   |        ^^^^^^
   |

也就是说,我们不能将带着类型的参数直接传给函数运行,而是需要将参数中的值取出来,再传给函数运行,即:

#[proc_macro_attribute]
pub fn deco(attr: TokenStream, item: TokenStream) -> TokenStream {
    let decorator = parse_macro_input!(attr as Ident);
    // ...
    let arguments: Vec<_> = inputs
	.iter()
	.map(|input| match input {
            FnArg::Typed(val) => &val.pat,
            _ => unreachable!()
        })
	.collect();

    let caller = quote! {
        #vis fn #ident(#inputs) #output {
            let func = #decorator(original_fn);
            return func(#(#arguments), *);

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

类型 FnArg 是一个枚举类型,包含 Typed 和 Received,前者表示普通的带有类型的参数,后者则专指 self 参数,因此在这里,我们只需要处理 Typed 参数。最后,将处理完毕的参数传给函数运行即可。完整代码如下:

use syn::{Ident, FnArg}

#[proc_macro_attribute]
pub fn deco(attr: TokenStream, item: TokenStream) -> TokenStream {
    let decorator = parse_macro_input!(attr as Ident);
    let decoratee = parse_macro_input!(item as ItemFn);

    let vis = &decoratee.vis;
    let ident = &decoratee.sig.ident;
    let block = &decoratee.block;
    let inputs = &decoratee.sig.inputs;
    let output = &decoratee.sig.output;

    let arguments: Vec<_> = inputs
	.iter()
	.map(|input| match input {
            FnArg::Typed(val) => &val.pat,
            _ => unreachable!()
        })
	.collect();

    let caller = quote! {
        #vis fn #ident(#inputs) #output {
            let func = #decorator(original_fn);
            return func(#(#arguments), *);

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

结语

当然,装饰器肯定不止这么一种,但我们可以根据装饰实现原理,借助 synquote 这两个神奇而强大的库,实现任意我们想要实现的装饰器。好啦,快到工作中用起来吧!