Rustでのデータ読み込みはSerdeIOにお任せ

概要

皆さんRust書いてますか？Rustの定番ライブラリにSerdeというものがあります。structにSerializeやDeserializeなどのマクロを追加するだけで様々な形式に書き出したり、逆に読み出したりできるようになります。

use serde::{Serialize, Deserialize};

#[derive(Debug, Deserialize, Serialize)]
struct User {
    id: u32,
    name: String,
    items: Vec<String>,
}

let user = serde_json::from_str(data)?;
let user_json_str = serde_json::to_string()?;

上記のコードでもわかるようにSerde自体はstructに情報を付加するだけであり、実際にJSONのencode/decodeはserde_jsonというcrateが行っています。今回はこのように特定のファイル形式に依存せずに汎用的に使用できるSerdeIOという自作のcrateをご紹介します。

API

SerdeIOのAPIは非常にシンプルで、以下の8つの関数のみです。

• read_record_from_reader
• read_records_from_reader
• read_record_from_file
• read_records_from_file
• write_record_to_writer
• write_records_to_writer
• write_record_to_file
• write_records_to_file

read/writeはそれぞれ読み込みと書き込みの関数を表します。
recordの場合はT、recordsの場合はVec<T>として扱うことを表します。
readerはio.Read、writerはio.Writeのtraitを実装した型に対してIOを行うことを表します。
fileの場合はファイルをopenしてそれに対してIOを行います。

例えばread_records_from_fileは実際には以下のようなシグニチャになっていてファイルからTの型のレコードを複数読み込んでVec<T>に格納して返す関数になります。

pub fn read_records_from_file<T: DeserializeOwned>(path: impl AsRef<Path>) -> AnyResult<Vec<T>>

具体例は以下の通りです。

use anyhow::{anyhow, Context, Result as AnyResult};
use serde::{Deserialize, Serialize};
use serdeio::{read_record_from_file, write_records_to_writer, DataFormat};

#[derive(Debug, Deserialize, Serialize)]
struct User {
    id: u32,
    name: String,
    items: Vec<String>,
}

pub fn main() -> AnyResult<()> {
    // get input file path from argv
    let args: Vec<String> = std::env::args().collect();
    let input_file_path = &args[1];

    // read json to memory
    let users: Vec<User> = read_record_from_file(input_file_path)
        .map_err(|e| anyhow! {e})
        .context("Failed to read records from file")?;

    // write to stdout in json lines format
    let writer = std::io::stdout();
    write_records_to_writer(writer, DataFormat::JsonLines, &users).unwrap();

    Ok(())
}

ファイルを開いてデータを読み込んでJSONをでコードしてというよくある流れを1つの関数で直接行えます。
現状以下の形式に対応しています。

• JSON
• JSON Lines
• CSV (csv crateを利用。csv featureを有効にする必要があります)
• YAML (serde_yaml crateを利用。yaml featureを有効にする必要があります)

フォーマットの制約により、JSON LinesとCSVはrecords系の関数のみに対応しています。私はよく計算系のアプリケーションを開発しています。開発時にはJSON Linesをよく使用するのですが。PoC時にお客さんがCSVでデータを渡してきてもコードをいじらずにそのまま対応することができます。

仕組み

最後にSerdeIOの仕組みについて簡単に説明します。基本的にはRustの強力な型推論を最大限利用しています。

pub fn read_records_from_reader<T: DeserializeOwned>(
    reader: impl Read,
    data_format: DataFormat,
) -> AnyResult<Vec<T>> {
    match data_format {
        DataFormat::Json => backend::json::read(reader),
        DataFormat::JsonLines => backend::jsonlines::read(reader),
        #[cfg(feature = "csv")]
        DataFormat::Csv => backend::csv::read(reader),
        #[cfg(feature = "yaml")]
        DataFormat::Yaml => backend::yaml::read(reader),
    }
}

pub fn read_records_from_file<T: DeserializeOwned>(path: impl AsRef<Path>) -> AnyResult<Vec<T>> {
    let (data_format, rdr) = open_file(path)?;
    read_records_from_reader(rdr, data_format)
}

read_records_from_fileはGenericsパラメータTを持ち、このTはSerdeのDeserializeがついている型になります。まず引数のファイルを開き、拡張子からフォーマットを判別してread_records_from_readerに渡します。例えば以下のように戻り値を受け取る変数に事前にVec<User>の型を宣言しておくことでTはUserで確定します。

let users: Vec<User> = read_records_from_file("users.json")?;

そしてread_records_from_readerではフォーマットに応じて実際のデコードの実装が呼ばれます。例えばJSON Linesの場合には以下のように一行ずつserde_json::from_strでデコードします。

pub fn read<T: DeserializeOwned>(reader: impl Read) -> AnyResult<Vec<T>> {
    let reader = BufReader::new(reader);
    let mut records: Vec<T> = Vec::new();
    for line in reader.lines() {
        let line = line?;
        let record: T = serde_json::from_str(&line)?;
        records.push(record);
    }
    Ok(records)
}

writeの場合も基本的に同じですが、引数に書き込むデータを取りますのでこれだけで型がすべて確定します。

pub fn write_records_to_file<T: Serialize>(
    path: impl AsRef<Path>,
    records: &[T],
) -> AnyResult<()> {
    let data_format = DataFormat::try_from(path.as_ref())?;
    let file = File::create(path)?;
    write_records_to_writer(file, data_format, records)
}

おまけ

同じ発想でPythonでもPydanticを用いたPydanticIOというライブラリも作っています。

ただし、PythonはRustと違って型推定がないのでreadの時には型を引数で与える必要があります。

from pydantic import BaseModel
from pydanticio import read_records_from_file, write_records_to_file

class User(BaseModel):
    name: str
    age: int


users = read_records_from_file("users.csv", User)
write_records_to_file("users.json", users)