pythonでcsvファイルの読み込みをちょっと便利に扱う方法を考えてみたりしてた
pandas使えば良いというのはあるけれど。csvファイルの読み込みにはcsv moduleを使う。
はじめに
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つの場合があるかもしれない。
これはもっとも原始的な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()))]) """