Implement Iterator protocol and resource management for CSV readers

- DefaultCSVReader and AthenaCSVReader now inherit from collections.abc.Iterator
- Added close() method to properly release file resources
- Added context manager support (__enter__/__exit__)
- Updated ResultSet to use reader's close() method
- Removed _csv_file instance variable from ResultSet (reader manages file lifecycle)
- No _closed flag needed - use _file is None check instead

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: laughingman7743

pyathena/s3fs/reader.py

@@ -2,10 +2,11 @@
 from __future__ import annotations
 
 import csv
-from typing import Any, Iterator, List, Optional, Tuple
+from collections.abc import Iterator
+from typing import Any, List, Optional, Tuple
 
 
-class DefaultCSVReader:
+class DefaultCSVReader(Iterator[List[str]]):
     """CSV reader using Python's standard csv module.
 
     This reader wraps Python's standard csv.reader and treats empty fields
@@ -33,9 +34,10 @@ def __init__(self, file_obj: Any, delimiter: str = ",") -> None:
             file_obj: File-like object to read from.
             delimiter: Field delimiter character.
         """
+        self._file: Optional[Any] = file_obj
         self._reader = csv.reader(file_obj, delimiter=delimiter)
 
-    def __iter__(self) -> Iterator[List[str]]:
+    def __iter__(self) -> "DefaultCSVReader":
         """Iterate over rows in the CSV file."""
         return self
 
@@ -46,17 +48,33 @@ def __next__(self) -> List[str]:
             List of field values as strings.
 
         Raises:
-            StopIteration: When end of file is reached.
+            StopIteration: When end of file is reached or reader is closed.
         """
+        if self._file is None:
+            raise StopIteration
         row = next(self._reader)
         # Python's csv.reader returns [] for empty lines; normalize to ['']
         # to represent a single empty field (consistent with single-value handling)
         if not row:
             return [""]
         return row
 
+    def close(self) -> None:
+        """Close the underlying file object."""
+        if self._file is not None:
+            self._file.close()
+            self._file = None
+
+    def __enter__(self) -> "DefaultCSVReader":
+        """Enter context manager."""
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit context manager and close resources."""
+        self.close()
+
 
-class AthenaCSVReader:
+class AthenaCSVReader(Iterator[List[Optional[str]]]):
     """CSV reader that distinguishes between NULL and empty string.
 
     This is the default reader for S3FSCursor.
@@ -87,10 +105,10 @@ def __init__(self, file_obj: Any, delimiter: str = ",") -> None:
             file_obj: File-like object to read from.
             delimiter: Field delimiter character.
         """
-        self._file = file_obj
+        self._file: Optional[Any] = file_obj
         self._delimiter = delimiter
 
-    def __iter__(self) -> Iterator[List[Optional[str]]]:
+    def __iter__(self) -> "AthenaCSVReader":
         """Iterate over rows in the CSV file."""
         return self
 
@@ -101,8 +119,10 @@ def __next__(self) -> List[Optional[str]]:
             List of field values, with None for NULL and '' for empty string.
 
         Raises:
-            StopIteration: When end of file is reached.
+            StopIteration: When end of file is reached or reader is closed.
         """
+        if self._file is None:
+            raise StopIteration
         line = self._file.readline()
         if not line:
             raise StopIteration
@@ -234,3 +254,17 @@ def _parse_unquoted_field(self, line: str, pos: int) -> Tuple[str, int]:
             pos += 1
 
         return value, pos
+
+    def close(self) -> None:
+        """Close the underlying file object."""
+        if self._file is not None:
+            self._file.close()
+            self._file = None
+
+    def __enter__(self) -> "AthenaCSVReader":
+        """Enter context manager."""
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit context manager and close resources."""
+        self.close()

pyathena/s3fs/result_set.py

@@ -80,7 +80,6 @@ def __init__(
         self._csv_reader_class: CSVReaderType = csv_reader or AthenaCSVReader
         self._fs = self._create_s3_file_system()
         self._csv_reader: Optional[Any] = None
-        self._csv_file: Optional[Any] = None
 
         if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location:
             self._init_csv_reader()
@@ -122,8 +121,8 @@ def _init_csv_reader(self) -> None:
         path = f"{bucket}/{key}"
 
         try:
-            self._csv_file = self._fs._open(path, mode="rb")
-            text_wrapper = TextIOWrapper(self._csv_file, encoding="utf-8")
+            csv_file = self._fs._open(path, mode="rb")
+            text_wrapper = TextIOWrapper(csv_file, encoding="utf-8")
 
             if self.output_location.endswith(".txt"):
                 # Tab-separated format (no header row)
@@ -228,8 +227,7 @@ def fetchall(
     def close(self) -> None:
         """Close the result set and release resources."""
         super().close()
-        if self._csv_file:
+        if self._csv_reader:
             with contextlib.suppress(Exception):
-                self._csv_file.close()
-            self._csv_file = None
-        self._csv_reader = None
+                self._csv_reader.close()
+            self._csv_reader = None