サクッと動かすRust

概要

Rustでのプログラミングは所有権などの複雑さから難しいものに思われがちです。この記事ではRustでのプログラムをサクッと実装して実行する方法についてご紹介します。例えばAWS Lambdaのハンドラや、Open AIのAPIを呼び出すなどの日常的な作業スクリプトを簡単に実装/実行できます。

前提

rust-analyzer拡張機能

まずVSCodeのrust-analyzer拡張機能をインストールしましょう。構文の静的解析によってコンパイルしなくても文法や型のエラーを検知することができます。また後述するようにGUI上でのテストランナー機能もあるのでそちらも便利です。

Rustでサクッと実装する

Result/Option

Rustでは処理中に例外が発生しうる値はResultというEnumにラップされます。

enum Result<T, E> {
    Ok(T),
    Err(E),
}

またnullになりうる値はOptionEnumにラップされます。

pub enum Option<T> {
    None,
    Some(T),
}

通常、こうしたEnumにラップされた値を利用するためにはmatch文などを使ったパターン分岐処理を行う必要があります。例えばファイルを文字列として読み込むstd::fs::read_to_string()はResult<String, Error>を返します。読み込んだ内容を利用するには以下のようにmatch文を使うなどして、ResultがOkである分岐で処理しなければいけません。

let result = std::fs::read_to_string("./test.tsv");

match result {
    Ok(content) => println!("{}", content),
    Err(e) => println!("{:?}", e),
}

Result/Optionは標準ライブラリを始めとしてRustのプログラム内では頻出します(例外とnullはありふれたものです)。それぞれについて分岐処理が強制されるのは例外やnullの処理を漏らさず行えるという点でRustの安全性を高めています。しかしながら、テストコードなど手元で実行する際など実行時エラーを許容できる場合もあります。そうした場合にも厳密に条件分岐を書くのはやや冗長に感じられるでしょう。

Result/Optionには値を取り出す際の分岐処理を省略するunwrap()という関数が実装されています。unwrap()を実行するとResultがErrだったりOptionがNoneだった場合は実行時エラーになります(Rustではpanicと呼ばれます)。

let result = std::fs::read_to_string("./test.tsv");
// resultがErrなら実行時エラー(panic)
println!("{}", result.unwrap());

また以下のようにメソッドチェーンによって後続処理を行う場合も便利です。メソッドチェーンの途中でResult/Optionが登場したら同様にunwrap()して処理を続けていくことができます。

let result = std::fs::read_to_string("./test.tsv");
// String -> &[u8]に変換し、0番目を取り出す
// get()の戻り値はOption型なのでさらにunwrap()して値を取り出す
let first_byte = result.unwrap().as_bytes().get(0).unwrap();

一方でデプロイしたWebサーバなどで実行時エラーになると困るので、プロダクションコードでunwrap()を使うのはやめておいた方がいいでしょう。テストや検証段階のコードでこうした記述を行い、後からmatch文を使った分岐処理に書き直すということを筆者はよくやります。

スマートポインタの受け渡し

ヒープ領域に値が格納される文字列を始めとしたスマートポインタには所有権があります。そうした変数を例えば関数の引数に渡すとその関数に所有権が移り、それ以降の処理で利用できなくなります。

fn print_text(text: String) {
    println!("{}", text);
}

let text = String::from("hoge");
// 所有権がprint_text()に移る
print_text(text);
// コンパイルエラー
println!("{}", text);

通常であれば所有権ルールに違反しないようにコードを修正する必要があります。

// println!()の第二引数以降は参照なのでこの関数の引数も参照でよかった
fn print_text(text: &String) {
    println!("{}", text);
}

let text = String::from("hoge");
// 参照を渡す
print_text(&text);
// 所有権が残っているのでエラーにならない
println!("{}", text);

一部のスマートポインタはclone()というメソッドを持っており、値を複製することができます(Cloneトレイトを実装しているスマートポインタがそれに当たります)。以下のように所有権を受け渡したくない箇所でclone()して値を複製してしまえば良い場面があります。

let text = String::from("hoge");
// clone()で文字列を複製してしまう
print_text(text.clone());
println!("{}", text);

ただしclone()は複製した値を格納するヒープ領域を新しく用意してそこにコピーするため、値の容量が大きいほどパフォーマンスに悪影響を及ぼします。unwrap()ほど重大な問題を引き起こさないとはいえプロダクションコードでは慎重に扱った方が良いでしょう。一方で検証用のコードでパフォーマンスを気にしない場合は所有権ルールへの対応を簡略化できるメリットがあります。

エラーの型定義

Result型を返す関数はエラーの型を指定する必要があります。単一の型のエラーのみを返す関数ならResultの型引数にそれを指定するだけですが、複数の型のエラーを返す場合は難しくなります。

例えば以下のようにファイルの読み込み時に発生するstd::io::Errorを返しうる関数があるとします。

fn read_file(path: String) -> Result<String, std::io::Error> {
    let result = std::fs::read_to_string(path)?;
    Ok(result)
}

この関数で追加で読み取った結果を数値型にパースする処理を入れたとします。String::parse::<u64>()はParseIntErrorを返すのでResult<u64, std::io::Error>返す定義のままだとコンパイルエラーになります。正しく実装するためにはstd::io::ErrorとParseIntErrorのいずれかを返すように関数の型定義を修正しなければいけません。

fn read_file() -> Result<u64, std::io::Error> {
    let result = std::fs::read_to_string(path)?;
    // コンパイルエラー
    let parsed = result.parse::<u64>()?;
    Ok(parsed)
}

複数の種類のエラーをまとめる型(いわゆるユニオン型)を定義するにはEnumを使います。Rustのエラー型はDebugとDisplayのトレイトを実装する必要があります。

use std::fmt::{Display, Formatter};

#[derive(Debug)]
enum ReadFileError {
    IoError(std::io::Error),
    ParseError(ParseIntError),
}

impl Display for ReadFileError {
    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
        match self {
            Self::IoError(e) => {
                write!(f, "{}", e)
            }
            Self::ParseError(e) => {
                write!(f, "{}", e)
            }
        }
    }
}

上記のエラーを返すように関数の型定義を修正します。このまま使うならそれぞれのResultをmatch文で分岐させて対応するエラーを返す実装になります。

fn read_file(path: &String) -> Result<u64, ReadFileError> {
    let result = match std::fs::read_to_string(path) {
        Ok(value) => value,
        Err(e) => return Err(ReadFileError::IoError(e)),
    };
    let parsed = match result.parse::<u64>() {
        Ok(value) => value,
        Err(e) => return Err(ReadFileError::ParseError(e)),
    };
    Ok(parsed)
}

それぞれのエラーが発生する場面でmatch文を書くのは面倒なので?で直接各エラーを返せるようにしたいところです。そのためには以下の自動変換を定義するFromトレイトを実装する必要があります。

impl From<std::io::Error> for ReadFileError {
    fn from(err: std::io::Error) -> Self {
        ReadFileError::IoError(err)
    }
}

impl From<ParseIntError> for ReadFileError {
    fn from(err: ParseIntError) -> Self {
        ReadFileError::ParseError(err)
    }
}

Fromトレイトを実装したことでReadFileErrorを返す関数をよりスマートに書くことができます。

fn read_file(path: &String) -> Result<u64, ReadFileError> {
    let result = std::fs::read_to_string(path)?;
    let parsed = result.parse::<u64>()?;
    Ok(parsed)
}

以上で見てきたように完全なエラーのEnumを実装するのはかなり手間がかかります。簡単な代替として戻り値のエラーをBox<dyn std::error::Error>と定義してしまう方法があります。

fn read_file(path: &String) -> Result<u64, Box<dyn std::error::Error>> {
    let result = std::fs::read_to_string(path)?;
    let parsed = result.parse::<u64>()?;
    Ok(parsed)
}

dynはトレイトオブジェクトを表しています。Boxは単純なスマートポインタで、トレイトオブジェクトのサイズが分からないとResultの型引数に入れられないので参照に変換しています。

Box<dyn std::error::Error>を返す関数を以下のように呼び出すことができます。

let path = String::from("./test.tsv");
match read_file(&path) {
    Ok(value) => println!("{}", value),
    Err(e) => println!("{:?}", e),
}

ここでのeは最も抽象的なエラートレイト(std::error::Error)を実装しているという以外の情報が失われています。そのためエラーの種類に応じて処理を分岐させることはできません。しかしながら、実行時エラーを許容できる環境ならunwrap()してしまえば何型だろうと関係なくその時点で処理が止まるので問題になりません。

let read_result = read_file(&path).unwrap();

Rustをサクッと実行する

通常、Rustはモジュールファイルでmain()を定義し、cargo runコマンドで実行します。

fn read_file() -> Result<u64, ReadFileError> {
    let result = std::fs::read_to_string("./test.tsv")?;
    let parsed = result.parse::<u64>()?;
    Ok(parsed)
}

fn main() {
    let path = String::from("./test.tsv");
    let content = match read_file(&path) {
        Ok(value) => value,
        Err(e) => {
            println!("{}", e);
            0
        }
    };
    println!("{}", content);
}

この場合実行できるエントリーポイントは一つだけです。処理を分岐させたいなら例えばコマンドライン引数で値を入力する必要があります。

Rustでは実行関数と同じファイルにテストモジュールを記述できます。#[cfg(test)]と書いておけばコンパイル対象に含まれません。

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test() {}
}

このテストコードを前節でご紹介した方法でサクッと実装して実行することでRustをスクリプト言語のように実行していくことができます。

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_file_1() {
        let path = String::from("./test_1.tsv");
        println!("{}", read_file(&path).unwrap());
    }

    #[test]
    fn test_file_2() {
        let path = String::from("./test_2.tsv");
        println!("{}", read_file(&path).unwrap());
    }
}

rust-analyzerをインストールしていると、各テスト関数の上に"Run Test"ボタンが表示されます。それをクリックするだけで各関数を実行することができます。Rustにはテストコードを個別に実行するコマンドがありそれを叩くユーティリティです。

cargo test \
    --package sample \
    --bin sample \
    -- sample::tests::test_file_1 \
    --exact --nocapture

これを利用して日常的な作業をRustでサクッと実行することができます。とはいえ処理する関数本体をプロダクション環境に耐えるように書き直した上で、assert!で結果の検証まで記述してしまえば立派な自動テストコードになります。

サクッと動かす例

以下では前節でご紹介したテクニックを使ってRustでサクッと実装/実行していく例をご紹介します。

Lambdaハンドラ

以下はlambda_runtimeクレートのサンプルコードです。数行でLambda関数のハンドラを実装することができます。

use lambda_runtime::{service_fn, LambdaEvent, Error};
use serde_json::{json, Value};

#[tokio::main]
async fn main() -> Result<(), Error> {
    let func = service_fn(func);
    lambda_runtime::run(func).await?;
    Ok(())
}

async fn func(event: LambdaEvent<Value>) -> Result<Value, Error> {
    let (event, _context) = event.into_parts();
    let first_name = event["firstName"].as_str().unwrap_or("world");

    Ok(json!({ "message": format!("Hello, {}!", first_name) }))
}

lambda_runtime - crates.io: Rust Package Registry

まずmain()やfunc()はlambda_runtime::Errorを返しています。このエラー型は実質的にBox<dyn std::error::Error>と同等のもので、どんなエラーでも含められます。つまりどんなResult型を返す関数を使っても?を使って中の値を取り出すことができます。

またfunc()ではunwrap_or()が使われています。unwrap()とほぼ同じですが、こちらはNoneの場合にpanicするのではなくデフォルトの値で置き換えます。

こうしたテクニックを使うことでひとまずLambda関数を動作させることができます。もちろんプロダクションコードにする際はそれぞれのエラーごとに例外処理をしたり、unwrap_or()するのではなく期待する値が入力されなかった場合の例外処理を行う方が良いでしょう。

Open AI APIを呼び出すスクリプト

筆者が実装したRustでOpen AIのCompletions APIを呼び出すスクリプトです。テストコードのtest_complete_text()を実行するとOpen AIのAPIにプロンプトが送信され、結果がコマンドラインに出力されます。complete_text()の引数のプロンプトを適当に変えて出力を検証することができます。

use reqwest::header::HeaderMap;
use serde::Serialize;
use std::env;

#[derive(Debug, Clone, Serialize)]
struct CompleteTextRequest {
// ...
}

fn main() {}

pub async fn complete_text(
    prompt: String,
) -> Result<serde_json::Value, Box<dyn std::error::Error>> {
    let client = reqwest::Client::new();
    let mut headers = HeaderMap::new();
    headers.append(
        "Authorization",
        format!("Bearer {}", env::var("OPENAI_API_KEY")?).parse()?,
    );
    let body = CompleteTextRequest {
        model: "text-davinci-003".to_string(),
        prompt: Some(prompt),
        max_tokens: Some(512),
        temperature: Some(0.7),
        top_p: Some(1),
        n: Some(1),
    };
    let response = client
        .post("https://api.openai.com/v1/completions")
        .headers(headers)
        .json(&body)
        .send()
        .await?;
    let response_body = response.json::<serde_json::Value>().await?;
    Ok(response_body)
}

#[cfg(test)]
mod tests {
    use super::*;
    use dotenv::dotenv;

    #[tokio::test]
    async fn test_complete_text() {
        dotenv().ok();
        let prompt = "say hello".to_string();
        let response = complete_text(prompt.clone()).await.unwrap();
        println!("prompt: {:?}", prompt);
        println!("response: {:?}", response);
    }
}

complete_text()は例えばenv::var("OPENAI_API_KEY")がVarErrorを、client.send()がrequest::Errorを返すので複数の型のエラーを返す関数です。ひとまずunwrap()する想定でBox<dyn std::error::Error>を返す実装にすることができます。しかしながら例外処理を正しく実装すればプロダクション環境でも利用できる関数になります。

プロンプトの文字列をひとまずclone()して関数に渡しています。プロンプトが長くなるとパフォーマンスに影響があるので、参照を渡すなど実装の修正を検討する必要があるでしょう。例のように短いプロンプトを検証するだけならこのままで問題ありません。

またAPIのレスポンスのJSONをserde_json::Valueを使ってパースしています。この状態では値を取り出す際に型が不明ですが、レスポンスの型定義を後で追加することもできます。

まとめ

Rustを使った実装を簡略化するためのいくつかのテクニックをご紹介しました。Rustの長所である安全性をいくらか損なう書き方であるため注意して扱う必要があります。一方で他のスクリプト言語にはないメリットとして、簡単に安全なコードに書き直すことができるという点があります。最初からRustで実装することでAPIの検証->プロダクションコードへの組み込みといった流れがスムーズです。

参考文献

Rust: Enums to wrap multiple errors

Wrapping errors - Rust By Example