Retrieving CSV fields and also raw string from file simultaneously in Python

Viewed 243

I have a generator that yields rows from a CSV file one at a time, something like:

import csv

def as_csv(filename):
    with open(filename) as fin:
        yield from csv.reader(fin)

However, I need to also capture the raw string returned from the file, as this needs to be persisted at the same time.

As far as I can tell, the csv built-in can be used on an ad-hoc basis, something like this:

import csv

def as_csv_and_raw(filename):
    with open(filename) as fin:
        for row in fin:
            raw = row.strip()
            values = csv.reader([raw])[0]
       yield (values, raw)

... but this has the overhead of creating a new reader and a new iterable for each row of the file, so on files with millions of rows I'm concerned about the performance impact.

It feels like I could create a coroutine that could interact with the primary function, yielding the parsed fields in a way where I can control the input directly without losing it, something like this:

import csv

def as_csv_and_raw(filename):
    with open(filename) as fin:
        reader = raw_to_csv(some_coroutine())
        reader.next()
        for row in fin:
            raw = row.strip()
            fields = reader.send(raw)
            yield fields, raw

def raw_to_csv(data):
    yield from csv.reader(data)

def some_coroutine():
    # what goes here?
    raise NotImplementedError

I haven't really wrapped my head around coroutines and using yield as an expression, so I'm not sure what goes in some_coroutine, but the intent is that each time I send a value in, that value is run through the csv.reader object and I get the set of fields back.

Can someone provide the implementation of some_coroutine, or alternately show me a better mechanism for getting the desired data?

1 Answers

You can use itertools.tee to create two independent iterators from the iterable file object, create a csv.reader from one of them, and then zip the other iterator with it for output:

from itertools import tee

def as_csv_and_raw(filename):
    with open(filename) as fin:
        row, raw = tee(fin)
        yield from zip(csv.reader(row), raw)
Related