pythonでcsvファイルの読み込みをちょっと便利に扱う方法を考えてみたりしてた

pandas使えば良いというのはあるけれど。csvファイルの読み込みにはcsv moduleを使う。

はじめに

対象のcsvは以下(csvと言いつつtsv)

target.tsv

hello    こんにちは
"hello world"   "こんにちは 世界"

以下の様なコードで読める

import csv

with open("./target.tsv") as rf:
    reader = csv.reader(rf, delimiter="\t")
    for row in reader:
        print(type(row), row)

"""
output:

<class 'list'> ['hello', 'こんにちは']
<class 'list'> ['hello world', 'こんにちは 世界']
"""

header付きのtsv

header付きのtsvを使いたい場合もあるかもしれない。

target2.tsv

word mean
hello   こんにちは
"hello world"   "こんにちは 世界"

とりあえず、next()で読み飛ばしてみる。

import csv

with open("./target2.tsv") as rf:
    reader = csv.reader(rf, delimiter="\t")
    next(reader)
    for row in reader:
        print(type(row), row)

DictReaderを使っても良いのでは?

せっかくheaderが付いているのにこれを利用しないのはもったいない。csv.DictReaderを使う事を考えてみる

import csv

with open("./target2.tsv") as rf:
    reader = csv.DictReader(rf, delimiter="\t")
    for row in reader:
        print(type(row), row)

"""
output:
<class 'dict'> {'mean': 'こんにちは', 'word': 'hello'}
<class 'dict'> {'mean': 'こんにちは 世界', 'word': 'hello world'}
"""

結果はdictで返ってくる。headerで指定したフィールド名でアクセスできるようになるが以前利用していた何番目の要素であるかをインデックスで指定するコードが流用できなくなる。

namedtupleを使ってみるのはどうだろう?

namedtupleを使ってみるのはどうだろう。これは事実上tupleのサブクラスなのだけれど、クラスを作成する際に利用した名前でもアクセス可能になる。

>>> from collections import namedtuple
>>> Point = namedtuple("Point", "x y")
>>> p = Point(x=1, y=2)
>>> p
Point(x=1, y=2)
>>> p = Point(*[1,2])
>>> p
Point(x=1, y=2)
>>> p.x
1
>>> p[0]
1

例えばheaderを常に取れるということを前提にすると以下の様な感じで書くこともできるかもしれない。

import csv
from collections import namedtuple


with open("./target2.tsv") as rf:
    reader = csv.reader(rf, delimiter="\t")
    Row = namedtuple("Row", " ".join(next(reader)))
    for line in reader:
        row = Row(*line)
        print(type(row), row)

"""
output:
<class '__main__.Row'> Row(word='hello', mean='こんにちは')
<class '__main__.Row'> Row(word='hello world', mean='こんにちは 世界')
"""

この場合には、namedtupleなのでインデックスでのアクセスも名前でのアクセスも有効になる。 (もっとも幾分かコストは掛かるようになるけれど)

クラス化してみる

namedtupleを使ってみるというのはそれなりに良いものではありそう。使い回しの効くようにクラス化しておくと便利かもしれない。

一方で今まではcsvファイルの先頭にヘッダー行があることを前提としていた。 本来は以下の2つの場合があるかもしれない。

  • csv中にヘッダー行が存在する
  • csv中にヘッダー行が存在しない。ある特定のフォーマットに従って値が列挙されている。

これはもっとも原始的なseriealizeの方法と言えるのかもしれない。 プロトコルとしてheader行にschemaを記述するという慣例があるというような感じ。 schemaの情報を事前に持っておけるようにもするべきとなると以下の様なコードになるかもしれない。

class WithSchemaReaderWrapper(object):
    def __init__(self, reader, schema=None, schema_factory=None, row_factory=None):
        self.schema = schema
        self.schema_factory = schema_factory
        self.reader = reader
        self.row_factory = row_factory or (lambda self, row: self.schema(*row))

    def create_row(self, row):
        return self.row_factory(self, row)

    def create_schema(self, row):
        return self.schema_factory(self, row)

    def __iter__(self):
        if self.schema is None:
            self.schema = self.create_schema(next(self.reader))
        return self

    def __next__(self):
        return self.create_row(next(self.reader))

これは以下の様にして使う。

import csv
from collections import namedtuple
from readerwrapper import WithSchemaReaderWrapper


def make_schema(_, row):
    return namedtuple("Row", " ".join(row))


with open("./target2.tsv") as rf:
    reader = csv.reader(rf, delimiter="\t")
    reader = WithSchemaReaderWrapper(reader, schema_factory=make_schema)
    for row in reader:
        print(type(row), row)

"""
output:
<class '__main__.Row'> Row(word='hello', mean='こんにちは')
<class '__main__.Row'> Row(word='hello world', mean='こんにちは 世界')
"""

事前にschemaを指定したい場合にはschemaに渡してあげれば良い。

import csv
from collections import namedtuple
from readerwrapper import WithSchemaReaderWrapper


with open("./target.tsv") as rf:
    reader = csv.reader(rf, delimiter="\t")
    Schema = namedtuple("Row", "name mean")
    reader = WithSchemaReaderWrapper(reader, schema=Schema)
    for row in reader:
        print(type(row), row)

"""
output:
<class '__main__.Row'> Row(name='hello', mean='こんにちは')
<class '__main__.Row'> Row(name='hello world', mean='こんにちは 世界')
"""

validation

schemaなのでvalidationをしたいと思うかもしれない。marshmallowあたりを使うと便利かもしれない。

pip install --user marshmallow

marshmallowのpost_loadを使うことでobjectへのmappingもできる

import csv
from marshmallow import Schema, fields, post_load
from readerwrapper import WithSchemaReaderWrapper


class WordSchema(Schema):
    name = fields.Str()
    mean = fields.Str()

    @post_load
    def make_instance(self, data):
        return Word(**data)


class Word(object):
    def __init__(self, name, mean):
        self.name = name
        self.mean = mean


def create_instance(self, data):
    return self.schema.load(data)


with open("./target.tsv") as rf:
    schema = WordSchema()
    reader = csv.DictReader(rf, delimiter="\t", fieldnames=list(schema.fields.keys()))
    reader = WithSchemaReaderWrapper(reader, schema=WordSchema(), row_factory=create_instance)
    for word, errors in reader:
        if errors:
            continue
        print(type(word), word, errors)

"""
output:
<class '__main__.Word'> <__main__.Word object at 0x1020457d0> {}
<class '__main__.Word'> <__main__.Word object at 0x10203a9d0> {}
"""

marshmallowにはmany=Trueというoptionがあるのでこれでも良いかもしれない。

今まで事前に定義しておいたWithSchemareaderWrapperを使っていたけれど、なくても良いのかもしれない。

import csv
from marshmallow import Schema, fields, validates, ValidationError


class WordSchema(Schema):

    name = fields.Str()
    mean = fields.Str()

    @validates("mean")
    def dont_include_white_space(self, s):
        if " " in s:
            raise ValidationError("oops")


def create_instance(self, data):
    return self.schema.load(data)


with open("./target.tsv") as rf:
    schema = WordSchema(many=True)
    reader = csv.DictReader(rf, delimiter="\t", fieldnames=list(schema.fields.keys()))
    words, errors = schema.load(reader)

    print(errors)
    print("-")
    for i, word in enumerate(words):
        if i in errors:
            continue
        print(i, type(word), word)

"""
output:

{1: {'mean': ['oops']}}
-
0 <class 'dict'> {'mean': 'hello', 'name': 'こんにちは'}
"""

エラーの情報も取れる。

もっともこのように考えるとdictをiterateするiteratorがあれば良いということになるのでcsvの読み込みはDictReaderで十分そう。

ちょっとお得なmarshmallowの使い方

marshmallowは入力された値によって良い感じに値をserializeしてくれる機能を持っていたりする。 デフォルトでサポートされている型は以下。

class BaseSchema(base.SchemaABC):

    TYPE_MAPPING = {
        text_type: fields.String,
        binary_type: fields.String,
        dt.datetime: fields.DateTime,
        float: fields.Float,
        bool: fields.Boolean,
        tuple: fields.Raw,
        list: fields.Raw,
        set: fields.Raw,
        int: fields.Integer,
        uuid.UUID: fields.UUID,
        dt.time: fields.Time,
        dt.date: fields.Date,
        dt.timedelta: fields.TimeDelta,
        decimal.Decimal: fields.Decimal,
    }

一度dumpを使うとその時に利用したfieldの型を覚えてくれるので何か便利な方法がある気がするけれど。 もう少し良い方法を考えたいところ。

target3.tsv

foo  20  2000-01-01T00:00:00Z
bar 10  2000-01-02T00:00:00Z

一度以下の様な形でschemaを作成してトレーニングしておく。

from marshmallow import Schema
from datetime import datetime


def schema_factory(sample, ordered=True):
    _ordered = ordered
    _fields = [p[0] for p in sample]

    class MySchema(Schema):
        class Meta:
            ordered = _ordered
            fields = _fields

    schema = MySchema()
    schema.dump(dict(sample))
    return schema


schema = schema_factory(sample=[("name", "foo"), ("age", 10), ("ctime", datetime.now())])

すると文字列だけからなるデータが適切なフォーマットに変換されて返される。 何かこの辺り便利な方法を探したい所ではある。

import csv

with open("./target3.tsv") as rf:
    reader = csv.DictReader(rf, delimiter="\t", fieldnames=schema.Meta.fields)
    for row in reader:
        result, errors = schema.load(row)
        # datetime,intなどに変換されている
        print(result)

"""
output:

OrderedDict([('name', 'foo'), ('age', 20), ('ctime', datetime.datetime(2000, 1, 1, 0, 0, tzinfo=tzutc()))])
OrderedDict([('name', 'bar'), ('age', 10), ('ctime', datetime.datetime(2000, 1, 2, 0, 0, tzinfo=tzutc()))])
"""