nomによるnumpyデータのパース

rust

nom

nomはrust製のパーサコンビネーターライブラリです． 個人的にはno_stdでの動作をサポートしてる点が嬉しいです．

今回勉強のためにnomでnumpyのファイル（のヘッダ）をパースしてみます．

numpyのフォーマット

今回は以下のようにして作成したデータをパースしてみます．

$ python -c "import numpy as np; a = np.arange(9).reshape(3,3); np.save('a.npy', a)";
$ hexdump -C a.npy
00000000  93 4e 55 4d 50 59 01 00  46 00 7b 27 64 65 73 63  |.NUMPY..F.{'desc|
00000010  72 27 3a 20 27 3c 69 38  27 2c 20 27 66 6f 72 74  |r': '<i8', 'fort|
00000020  72 61 6e 5f 6f 72 64 65  72 27 3a 20 46 61 6c 73  |ran_order': Fals|
00000030  65 2c 20 27 73 68 61 70  65 27 3a 20 28 33 2c 20  |e, 'shape': (3, |
00000040  33 29 2c 20 7d 20 20 20  20 20 20 20 20 20 20 0a  |3), }          .|
00000050  00 00 00 00 00 00 00 00  01 00 00 00 00 00 00 00  |................|
00000060  02 00 00 00 00 00 00 00  03 00 00 00 00 00 00 00  |................|
00000070  04 00 00 00 00 00 00 00  05 00 00 00 00 00 00 00  |................|
00000080  06 00 00 00 00 00 00 00  07 00 00 00 00 00 00 00  |................|
00000090  08 00 00 00 00 00 00 00                           |........|
00000098

フォーマットの公式ドキュメント: https://docs.scipy.org/doc/numpy-dev/neps/npy-format.html

以下のようなフォーマットになっているようです．

  • 先頭は \x93NUMPY
  • 次の1byteがメジャーバージョン
  • 次の1byteがマイナーバージョン
  • メジャーバージョン1の場合，次の2byteがヘッダ長（この後に続くpythonのdictを表す文字列の長さ（パディング含む））
  • メジャーバージョン2の場合，次の4byteがヘッダ長
  • 次に文字列でpythonのdictの形式でdescr, fortran_order, shape が格納されている
    • descr : numpyのdtype．通常は'<u8'とか'<f8'などの文字列だが，structured arrayの場合は型情報のリストになる．
    • fortran_order : TrueかFalse
    • shape: タプルでarrayの形を表す
  • その後，アラインメントを揃えるために適当なスペースがあり，最後に改行(\x0a)
  • その後が実際のデータ．データの格納形式はdescrで指定されたdtypeに基づく

完成品1

とりあえず，1. dictの部分はdescr, fortran_order, shapeの順に格納されていると仮定 *1 2. dtypeの部分はstrucured arrayは考えないで，endianとワードサイズのみだけ取り出す として作ってみました．

#[derive(Debug, PartialEq)]
pub struct NpyHeader {
    major_version: u8,
    minor_version: u8,
    header_len: u32,
    little_endian: bool,
    fortran_order: bool,
    word_size: u32,
    shape: Vec<u8>,
}
named!(pub parse_header<NpyHeader>,
    do_parse!(
        tag!(b"\x93NUMPY") >>
        major_version: le_u8 >>
        minor_version: le_u8 >>
        header_len: alt!(
            cond_reduce!(major_version == 1, map!(le_u16, |x| x as u32)) |
            cond_reduce!(major_version == 2, le_u32)) >>
        tag!("{'descr': '") >>
        little_endian: map!(alt!(tag!("<") | tag!(">")), |x| x == b"<") >>
        le_u8 >> // skip byte
        word_size: map_res!(map_res!(digit, std::str::from_utf8), std::str::FromStr::from_str) >>
        tag!("', 'fortran_order': ") >>
        fortran_order: map!(alt!(tag!("True") | tag!("False")), |s| s == b"True") >>
        tag!(", 'shape': ") >>
        shape: delimited!(char!('('),
                          separated_list!(ws!(char!(',')),
                                            map_res!(
                                                map_res!(digit, std::str::from_utf8),
                                                    std::str::FromStr::from_str)),
                          char!(')')) >>
        tag!(", }") >>
        take_while!(call!(|c| c == b' ')) >>
        tag!("\n") >>
        (
            NpyHeader{
                major_version,
                minor_version,
                header_len,
                little_endian,
                fortran_order,
                word_size,
                shape,
            }
        )
    )
);

実行結果

#[cfg(test)]
mod test {
    use super::*;
    use std::fs::File;
    use std::io::Read;

    #[test]
    fn parse_nom() {
        let mut buf = vec![];
        File::open("./a.npy")
            .expect("failed to open file")
            .read_to_end(&mut buf)
            .unwrap();
        let r = parse_header(&buf);
        println!("{:?}", r);
        assert_eq!(
            r.to_result().ok().unwrap(),
            NpyHeader {
                major_version: 1,
                minor_version: 0,
                header_len: 70,
                little_endian: true,
                fortran_order: false,
                word_size: 8,
                shape: vec![3, 3],
            }
        );
    }
}

$ cargo test -- --nocapture
...
running 1 test
Done([0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0, 0, 0], NpyHeader { major_version: 1, minor_version: 0, header_len: 70, little_endian: true, fortran_order: false, word_size: 8, shape: [3, 3] })
test test::parse_nom ... ok

ポイント

  • named!() マクロで関数を定義する．今回の場合 fn parse_header(input: &[u8]) -> nom::IResult<&[u8], NpyHeader> みたいな関数が定義されることになる．
    • nom::IResult<I,O>の型パラメータは，Iがまだパースしていない残りの部分の型，Oがパースした結果の型
    • 正しくパースされると，Done(I,O)が返る．エラーだったらError(Err)かIncomplete(Needed)
  • do_parse!()マクロを使うと，IResultを返す関数を>>で繋げることができる
    • 具体的には，do_parse!(I->IResult<I,A> >> I->IResult<I,B> >> ... I->IResult<I,X> , ( O ) ) => I -> IResult<I, O> の形になる
    • 最後に()で囲んだものが最終的な戻り値
    • 基本的にnomがいろいろな関数とマクロを提供しているので，それをdo_parse!()の中で組み合わせていくことになる
    • do_parse!()の中でxxx: yyy とするとパース処理した結果が後からxxxで参照できる
  • tag!()で指定したバイト列の読み取り
  • le_u8はlittle endianで1byte読み取る関数*2．同様の関数にle_u16とかle_u32など
  • alt!() は複数のパーサを受け取り，先頭から試していって最初に成功した結果を返す
    • alt!(tag!("True") | tag!("False")の部分は最初がtag!("True")（短い方の文字列)でなければならない．これは"False"が先の場合，"False"にマッチするか確認するために5文字読み込むので，次に"True"にマッチするか確認する際は文字長が足らないため必ず失敗する．詳細はドキュメント参照
  • パース処理で得られた部分の型を変換したい場合は，map!()あるいはmap_res!()を使う
    • スライスを文字列へ変換: map_res!(xxxx, std::str::from_utf8)
    • 文字列を数値へ変換: map_res!(xxxx, std::str::FromStr::from_str)
  • if処理をしたいときはcond!(), if-else処理をしたい時はalt!()とcond_reduce!()を組み合わせる
  • delimited!(opening, X, closing)でXを取り出す
  • ws!()は各トークン間のスペースを自動で取り除く
  • separated_list!()でセパレータで区切られた値をVecで取得できる．ちなみに，もし末尾にセパレータがある場合そのセパレータは残る (例えば1,2,3,みたいな場合)

完成品2

もう少し真面目にdictをパースしてみたのがこちら．といっても手抜きでdictのvalueとしてboolかstringか数値のタプルかの3つしか考えてないですが．上の例はあえて全部一つのdo_parse!()にまとめましたが実際には分割した方が，可読性もテストしやすさも向上します．本当ならdictをパースしたあと，さらに個別の各要素をチェックしていく必要がありますが，飽きてきたのでこの辺で.. ちなみにもしdict部分に変なデータがあってもそのままパースされることになります．

named!(to_u8<u8>,
    map_res!(
        map_res!(digit, std::str::from_utf8),
        std::str::FromStr::from_str)
);

#[derive(Debug, PartialEq)]
pub enum DictValue<'a> {
    Str(&'a str),
    Bool(bool),
    Tuple(Vec<u8>),
}

named!(tuple_value<DictValue>,
    map!(
        delimited!(
            char!('('),
            separated_list!(ws!(char!(',')), to_u8),
            pair!(opt!(ws!(char!(','))), char!(')'))),
        DictValue::Tuple)
);

named!(bool_value<DictValue>,
    map!(alt!(tag!("True") | tag!("False")),
        |s| DictValue::Bool(s == b"True"))
);

named!(quoted_string<&str>,
    map_res!(
        delimited!(
            char!('\''),
            is_not!("'"),
            char!('\'')),
        std::str::from_utf8)
);

named!(string_value<DictValue<'a>>,
    map!(
        map_res!(
            delimited!(
                char!('\''),
                is_not!("'"),
                char!('\'')),
            std::str::from_utf8),
        DictValue::Str)
);

named!(key_value<(&str, DictValue)>,
    do_parse!(
        key: quoted_string >>
        ws!(char!(':')) >>
        value: alt!(string_value | bool_value | tuple_value) >>
        (key, value)
    )
);

named!(pub parse_dict<HashMap<&str,DictValue>>,
    delimited!(
        pair!(char!('{'), opt!(multispace)),
        map!(many0!(terminated!(key_value, opt!(ws!(char!(','))))),
                |vec: Vec<_>| vec.into_iter().collect()),
        pair!(opt!(multispace), char!('}')))
);

#[derive(Debug, PartialEq)]
pub struct NpyHeader2<'a> {
    major_version: u8,
    minor_version: u8,
    header_len: u32,
    dict: HashMap<&'a str, DictValue<'a>>,
}

named!(pub parse_header2<NpyHeader2>,
    do_parse!(
        tag!(b"\x93NUMPY") >>
        major_version: le_u8 >>
        minor_version: le_u8 >>
        header_len: alt!(
            cond_reduce!(major_version == 1, map!(le_u16, |x| x as u32)) |
            cond_reduce!(major_version == 2, le_u32)) >>
        dict: parse_dict >>
        take_while!(call!(|c| c == b' ')) >>
        tag!("\n") >>
        (
            NpyHeader2{
                major_version,
                minor_version,
                header_len,
                dict,
            }
        )
    )
);

実行結果

    #[test]
    fn parse_nom2() {
        let mut buf = vec![];
        File::open("./a.npy")
            .expect("failed to open file")
            .read_to_end(&mut buf)
            .unwrap();
        let r = parse_header2(&buf);
        let mut dict = HashMap::new();
        dict.insert("descr", DictValue::Str("<i8"));
        dict.insert("fortran_order", DictValue::Bool(false));
        dict.insert("shape", DictValue::Tuple(vec![3,3]));
        println!("{:?}", r);
        assert_eq!(
            r.to_result().ok().unwrap(),
            NpyHeader2 {
                major_version: 1,
                minor_version: 0,
                header_len: 70,
                dict: dict,
            }
        );
    }

$ cargo test -- --nocapture
...
Done([0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0, 0, 0], NpyHeader2 { major_version: 1, minor_version: 0, header_len: 70, dict: {"descr": Str("<i8"), "fortran_order": Bool(false), "shape": Tuple([3, 3])} })

ポイント

  • alt!()を使う場合，各パーサの戻り値は同じでなければならない．今回はenumを使って対処している

余談

後から気づきましたが，nomを使ってnumpyのstructured arraysをパースしてるnpy-rsというのがあるのでnumpyデータをちゃんとパースしたい人は参考になるかもしれません．dictのパースはValueのenumを作ってパースしてけば綺麗に書ける訳ですね．確かに．

感想

マクロはあまり好きではないので最初とっつきにくい印象がありましたが，書いていくうちにまぁこんなもんかもしれないと思いました．同じ処理をするにもいろいろな書き方ができるので，ライブラリが提供しているマクロ/関数をうまく活用することが簡潔なコードを書く鍵になります．幸いnomで書かれたパーサはいろいろあるのでそれが参考になります．書いていて少し気になったのが，synstasticでの文法チェックが少々遅いこと．体感として明らかにnomでマクロを定義する前と後で遅くなりました．今回の簡単な例でも1s近く遅くなった気がします．これはnom（というかマクロ）の問題なのか，それ以外の問題なのかはわかりませんが..

*1:多分必ずそうなる保証はない

*2:1byteにendianも何もないですが