From b7905ff441d996178d1679d3b1e5cffd6d10ae78 Mon Sep 17 00:00:00 2001 From: "junjie.jiang" Date: Fri, 5 May 2023 17:01:43 +0800 Subject: [PATCH] Add csv reader Signed-off-by: junjie.jiang --- __init__.py | 4 ++++ read_csv.py | 20 ++++++++++++++++++++ 2 files changed, 24 insertions(+) create mode 100644 __init__.py create mode 100644 read_csv.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..03b200e --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ +from .read_csv import ReadCSV + +def csv_reader(*args, **kwargs): + ReadCSV(*args, **kwargs) diff --git a/read_csv.py b/read_csv.py new file mode 100644 index 0000000..9a394bc --- /dev/null +++ b/read_csv.py @@ -0,0 +1,20 @@ +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