update files

Author: ramanishsingh

test/nodes/test_csv_reader.py

@@ -105,7 +105,9 @@ def test_load_wrong_state(self):
 
         state = node.state_dict()
         state[CSVReader.HEADER_KEY] = None
-        with self.assertRaisesRegex(ValueError, "Check if has_header=True matches the state header=None"):
+        with self.assertRaisesRegex(
+            ValueError, "Check if has_header=True matches the state header=None"
+        ):
             node.reset(state)
 
         node.close()
@@ -133,7 +135,9 @@ def test_empty_file(self):
         node.close()
 
     def test_header_validation(self):
-        with self.assertRaisesRegex(ValueError, "return_dict=True requires has_header=True"):
+        with self.assertRaisesRegex(
+            ValueError, "return_dict=True requires has_header=True"
+        ):
             CSVReader("dummy.csv", has_header=False, return_dict=True)
 
     def test_multi_epoch(self):

torchdata/nodes/csv_reader.py

@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 import csv
 from itertools import islice
 from typing import Any, Dict, Iterator, List, Optional, Sequence, TextIO, Union
@@ -54,16 +60,22 @@ def reset(self, initial_state: Optional[Dict[str, Any]] = None):
     def _handle_initial_state(self, state: Dict[str, Any]):
         """Restore reader state from checkpoint."""
         # Validate header compatibility
-        if (not self.has_header and self.HEADER_KEY in state) or (self.has_header and state[self.HEADER_KEY] is None):
-            raise ValueError(f"Check if has_header={self.has_header} matches the state header={state[self.HEADER_KEY]}")
+        if (not self.has_header and self.HEADER_KEY in state) or (
+            self.has_header and state[self.HEADER_KEY] is None
+        ):
+            raise ValueError(
+                f"Check if has_header={self.has_header} matches the state header={state[self.HEADER_KEY]}"
+            )
 
         self._header = state.get(self.HEADER_KEY)
         target_line_num = state[self.NUM_LINES_YIELDED]
         assert self._file is not None
         # Create appropriate reader
         if self.return_dict:
 
-            self._reader = csv.DictReader(self._file, delimiter=self.delimiter, fieldnames=self._header)
+            self._reader = csv.DictReader(
+                self._file, delimiter=self.delimiter, fieldnames=self._header
+            )
         else:
             self._reader = csv.reader(self._file, delimiter=self.delimiter)
         # Skip header if needed (applies only when file has header)
@@ -112,5 +124,5 @@ def get_state(self) -> Dict[str, Any]:
         }
 
     def close(self):
-        if self._file and not self._file.closed:
+        if self._file is not None and not self._file.closed:
             self._file.close()