diff --git a/task-sdk/src/airflow/sdk/io/path.py b/task-sdk/src/airflow/sdk/io/path.py index 5b66e09d110a0..611dfe2c8df1b 100644 --- a/task-sdk/src/airflow/sdk/io/path.py +++ b/task-sdk/src/airflow/sdk/io/path.py @@ -60,6 +60,11 @@ def wrapper(*args, **kwargs): return wrapper + # We need to explicitly implement `__iter__` + # because otherwise python iteration logic could use `__getitem__` + def __iter__(self): + return iter(self._obj) + def __getitem__(self, key): # Intercept item access return self._obj[key] diff --git a/task-sdk/tests/task_sdk/io/test_path.py b/task-sdk/tests/task_sdk/io/test_path.py index 78c75cebddf6e..99a513664945c 100644 --- a/task-sdk/tests/task_sdk/io/test_path.py +++ b/task-sdk/tests/task_sdk/io/test_path.py @@ -215,6 +215,15 @@ def test_read_write(self, target): assert o.open("rb").read() == b"foo" o.unlink() + def test_read_line_by_line(self, target): + o = ObjectStoragePath(f"file://{target}") + with o.open("wb") as f: + f.write(b"foo\nbar\n") + with o.open("rb") as f: + lines = list(f) + assert lines == [b"foo\n", b"bar\n"] + o.unlink() + def test_stat(self, target): o = ObjectStoragePath(f"file://{target}") assert o.stat().st_size == 0