csv-reader/read_csv.py

import csv
from towhee.operator import PyOperator


class ReadCSV(PyOperator):
    """
    Wrapper of python csv: https://docs.python.org/3.8/library/csv.html
    """
    def __init__(self, f_path: str, newline='', dialect='excel', **fmtparams):
        self._f_path = f_path
        self._newline = newline
        self._dialect = dialect
        self._fmtparams = fmtparams

    def __call__(self):
        with open(self._f_path, newline=self._newline) as f:
            reader = csv.reader(f, self._dialect, **self._fmtparams)
            # skip header
            next(reader)
            yield from reader