This is an automated email from the ASF dual-hosted git repository.
blue pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iceberg.git
The following commit(s) were added to refs/heads/master by this push:
new 64a21b5ce8 Python: Ensure that the streams use context managers (#5436)
64a21b5ce8 is described below
commit 64a21b5ce8a8cdb1d841dcc7669544097d6ba43f
Author: Fokko Driesprong <[email protected]>
AuthorDate: Sun Aug 7 18:53:46 2022 +0200
Python: Ensure that the streams use context managers (#5436)
---
python/pyiceberg/io/base.py | 16 ++++++++++++++++
python/pyiceberg/io/memory.py | 6 ++++++
python/pyiceberg/serializers.py | 8 ++++----
python/tests/avro/test_decoder.py | 6 ++++++
4 files changed, 32 insertions(+), 4 deletions(-)
diff --git a/python/pyiceberg/io/base.py b/python/pyiceberg/io/base.py
index 3098102b99..6169dcdfd5 100644
--- a/python/pyiceberg/io/base.py
+++ b/python/pyiceberg/io/base.py
@@ -51,6 +51,14 @@ class InputStream(Protocol):
def close(self) -> None:
...
+ @abstractmethod
+ def __enter__(self):
+ ...
+
+ @abstractmethod
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ ...
+
@runtime_checkable
class OutputStream(Protocol): # pragma: no cover
@@ -68,6 +76,14 @@ class OutputStream(Protocol): # pragma: no cover
def close(self) -> None:
...
+ @abstractmethod
+ def __enter__(self):
+ ...
+
+ @abstractmethod
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ ...
+
class InputFile(ABC):
"""A base class for InputFile implementations
diff --git a/python/pyiceberg/io/memory.py b/python/pyiceberg/io/memory.py
index 73f0d5125d..28db68c1f0 100644
--- a/python/pyiceberg/io/memory.py
+++ b/python/pyiceberg/io/memory.py
@@ -70,3 +70,9 @@ class MemoryInputStream(InputStream):
def close(self) -> None:
del self.buffer
self.pos = 0
+
+ def __enter__(self):
+ pass
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self.close()
diff --git a/python/pyiceberg/serializers.py b/python/pyiceberg/serializers.py
index eee0c435ac..d922cf7944 100644
--- a/python/pyiceberg/serializers.py
+++ b/python/pyiceberg/serializers.py
@@ -54,7 +54,8 @@ class FromInputFile:
TableMetadata: A table metadata instance
"""
- return FromByteStream.table_metadata(byte_stream=input_file.open(),
encoding=encoding)
+ with input_file.open() as input_stream:
+ return FromByteStream.table_metadata(byte_stream=input_stream,
encoding=encoding)
class ToOutputFile:
@@ -70,6 +71,5 @@ class ToOutputFile:
output_file (OutputFile): A custom implementation of the
iceberg.io.file.OutputFile abstract base class
overwrite (bool): Where to overwrite the file if it already
exists. Defaults to `False`.
"""
- f = output_file.create(overwrite=overwrite)
- f.write(metadata.json().encode("utf-8"))
- f.close()
+ with output_file.create(overwrite=overwrite) as output_stream:
+ output_stream.write(metadata.json().encode("utf-8"))
diff --git a/python/tests/avro/test_decoder.py
b/python/tests/avro/test_decoder.py
index 0616d49314..15b3029ec1 100644
--- a/python/tests/avro/test_decoder.py
+++ b/python/tests/avro/test_decoder.py
@@ -119,6 +119,12 @@ class OneByteAtATimeInputStream(InputStream):
def close(self) -> None:
pass
+ def __enter__(self):
+ pass
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self.close()
+
def test_read_single_byte_at_the_time():
decoder = BinaryDecoder(OneByteAtATimeInputStream())
