tostorchconnector 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

tostorchconnector/tos_object_reader.py

@@ -1,5 +1,6 @@
 import io
 import logging
+import threading
 from functools import cached_property
 from os import SEEK_SET, SEEK_CUR, SEEK_END
 from typing import Optional, Callable, Any
@@ -28,6 +29,8 @@ class TosObjectReader(io.BufferedIOBase):
         self._total_size: Optional[int] = None
         self._read_offset = 0
         self._buffer = io.BytesIO()
+        self._closed = False
+        self._lock = threading.Lock()
 
     @property
     def bucket(self) -> str:
@@ -37,6 +40,26 @@ class TosObjectReader(io.BufferedIOBase):
     def key(self) -> str:
         return self._key
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is not None:
+            try:
+                log.info(f'Exception occurred before closing stream: {exc_type.__name__}: {exc_val}')
+            except:
+                pass
+        else:
+            self.close()
+
+    def close(self) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            if self._object_stream and isinstance(self._object_stream, ReadStream):
+                self._object_stream.close()
+
     def read(self, size: Optional[int] = None) -> Optional[bytes]:
         if self._is_read_to_end():
             return b''
@@ -52,11 +75,14 @@ class TosObjectReader(io.BufferedIOBase):
                     while 1:
                         chunk = self._object_stream.read(self._object_stream_offset, chunk_size)
                         if not chunk:
+                            self.close()
                             break
                         self._object_stream_offset += len(chunk)
                         self._buffer.write(chunk)
-                finally:
-                    self._object_stream.close()
+                except:
+                    self.close()
+                    raise
+
             elif isinstance(self._object_stream, GetObjectOutput):
                 for chunk in self._object_stream:
                     self._buffer.write(chunk)
@@ -149,13 +175,15 @@ class TosObjectReader(io.BufferedIOBase):
                 while offset > size:
                     chunk = self._object_stream.read(self._object_stream_offset, chunk_size)
                     if not chunk:
+                        self.close()
                         break
                     size += self._buffer.write(chunk)
                     self._object_stream_offset += len(chunk)
                 self._total_size = self._buffer.tell()
-            finally:
-                self._object_stream.close()
-        else:
+            except:
+                self.close()
+                raise
+        elif isinstance(self._object_stream, GetObjectOutput):
             try:
                 while offset > size:
                     size += self._buffer.write(next(self._object_stream))

{tostorchconnector-1.0.2.dist-info → tostorchconnector-1.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tostorchconnector
-Version: 1.0.2
+Version: 1.0.4
 Summary: TOS connector integration for PyTorch
 Author-email: xiangshijian <xiangshijian@bytedance.com>
 Classifier: Development Status :: 4 - Beta

{tostorchconnector-1.0.2.dist-info → tostorchconnector-1.0.4.dist-info}/RECORD

@@ -5,10 +5,10 @@ tostorchconnector/tos_common.py,sha256=oxxxzpA7Adn3hwNkpB9fHCueGGmZEc2f9tFqtEgPn
 tostorchconnector/tos_iterable_dataset.py,sha256=a28XPfZJwvWJJeI3qhHXJAAhKyg4U1DHavUI3VsQorY,4312
 tostorchconnector/tos_map_dataset.py,sha256=o8Fq8O5aG2LS8NA9eFtEfBbkisWwHg-dItO8UAqu08o,3471
 tostorchconnector/tos_object_meta.py,sha256=YrEQikioD5v0C_VcoudvTn0apUcxqxmSNQ4dDIip1Zc,562
-tostorchconnector/tos_object_reader.py,sha256=G6TA3IweJl1BzECmbYibZS7trPUq37Uxlz0EGS9Ul4Y,5647
+tostorchconnector/tos_object_reader.py,sha256=g2Iek4eDQllbr5AeFe81kB_zKSBBw18BGYElZqbnPcE,6472
 tostorchconnector/tos_object_writer.py,sha256=-I2coKCuUITL_LHf5we3cZBQ6rdn0bQL6NwpF39uDuU,1949
-tostorchconnector-1.0.2.dist-info/licenses/LICENSE,sha256=WBJyZyF8q8ZxohuLttIbv7HNLBRA0NnnRWWgfDlDZBE,11361
-tostorchconnector-1.0.2.dist-info/METADATA,sha256=G2BM2Q1ECH6rC8s7rkmHbqmSMzb2QZ95hRPGb3PZ5Lk,881
-tostorchconnector-1.0.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tostorchconnector-1.0.2.dist-info/top_level.txt,sha256=w9WvBP6KEi5Dourf86IdBbpWe-fad84S4SBftSg7H-k,18
-tostorchconnector-1.0.2.dist-info/RECORD,,
+tostorchconnector-1.0.4.dist-info/licenses/LICENSE,sha256=WBJyZyF8q8ZxohuLttIbv7HNLBRA0NnnRWWgfDlDZBE,11361
+tostorchconnector-1.0.4.dist-info/METADATA,sha256=wQTa2P1xE5nSaoyHGABTHHRF1lqx29mQeD6OyoXWnE0,881
+tostorchconnector-1.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tostorchconnector-1.0.4.dist-info/top_level.txt,sha256=w9WvBP6KEi5Dourf86IdBbpWe-fad84S4SBftSg7H-k,18
+tostorchconnector-1.0.4.dist-info/RECORD,,