diff --git a/clients/python-wrapper/lakefs/object.py b/clients/python-wrapper/lakefs/object.py index 3b12e82ab22..05146f5d48f 100644 --- a/clients/python-wrapper/lakefs/object.py +++ b/clients/python-wrapper/lakefs/object.py @@ -146,7 +146,10 @@ def writelines(self, lines: List[AnyStr]) -> None: raise io.UnsupportedOperation def __next__(self) -> AnyStr: - return self.readline() + line = self.readline() + if len(line) == 0: + raise StopIteration + return line def __iter__(self) -> Iterator[AnyStr]: return self diff --git a/clients/python-wrapper/tests/integration/test_object.py b/clients/python-wrapper/tests/integration/test_object.py index 5969abe7d17..aa9057252f0 100644 --- a/clients/python-wrapper/tests/integration/test_object.py +++ b/clients/python-wrapper/tests/integration/test_object.py @@ -306,3 +306,26 @@ def test_readline_partial_line_buffer(setup_repo): assert read == expected assert reader.read() == "" + + +def test_write_read_csv(setup_repo): + _, repo = setup_repo + columns = ["ID", "Name", "Email"] + sample_data = [ + ['1', "Alice", "alice@example.com"], + ['2', "Bob", "bob@example.com"], + ['3', "Carol", "carol@example.com"], + ] + obj = repo.branch("main").object(path="csv/sample_data.csv") + + with obj.writer(mode='w', pre_sign=False, content_type="text/csv") as fd: + writer = csv.writer(fd) + writer.writerow(columns) + for row in sample_data: + writer.writerow(row) + + for i, row in enumerate(csv.reader(obj.reader(mode='r'))): + if i == 0: + assert row == columns + else: + assert row == sample_data[i - 1]