From 4402af8a02f52dcee723dc1f5f0c3d1053b60c87 Mon Sep 17 00:00:00 2001 From: Michael Penkov Date: Thu, 21 Mar 2024 15:27:59 +0900 Subject: [PATCH 1/3] add support for zst writing --- smart_open/compression.py | 4 ++-- smart_open/tests/test_compression.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/smart_open/compression.py b/smart_open/compression.py index 499f557d..1c88224e 100644 --- a/smart_open/compression.py +++ b/smart_open/compression.py @@ -106,8 +106,8 @@ def _handle_gzip(file_obj, mode): def _handle_zstd(file_obj, mode): - import zstandard as zstd - result = zstd.ZstdDecompressor().stream_reader(file_obj, closefd=True) + import zstandard # type: ignore + result = zstandard.open(filename=file_obj, mode=mode) return result diff --git a/smart_open/tests/test_compression.py b/smart_open/tests/test_compression.py index 56ec100e..0675086e 100644 --- a/smart_open/tests/test_compression.py +++ b/smart_open/tests/test_compression.py @@ -7,6 +7,7 @@ # import gzip import io +import tempfile import pytest import zstandard as zstd @@ -42,3 +43,15 @@ def label(thing, name): def test_compression_wrapper_read(fileobj, compression, filename): wrapped = smart_open.compression.compression_wrapper(fileobj, 'rb', compression, filename) assert wrapped.read() == plain + + +def test_zst_write(): + with tempfile.NamedTemporaryFile(suffix=".zst") as tmp: + with smart_open.open(tmp.name, "wt") as fout: + print("hello world", file=fout) + print("this is a test", file=fout) + + with smart_open.open(tmp.name, "rt") as fin: + got = list(fin) + + assert got == ["hello world\n", "this is a test\n"] From 4a450e140d060f67904b8a8b077d4777df4c7c43 Mon Sep 17 00:00:00 2001 From: Michael Penkov Date: Thu, 21 Mar 2024 15:48:16 +0900 Subject: [PATCH 2/3] windows dammit --- smart_open/tests/test_compression.py | 13 ------------- smart_open/tests/test_smart_open.py | 12 ++++++++++++ 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/smart_open/tests/test_compression.py b/smart_open/tests/test_compression.py index 0675086e..56ec100e 100644 --- a/smart_open/tests/test_compression.py +++ b/smart_open/tests/test_compression.py @@ -7,7 +7,6 @@ # import gzip import io -import tempfile import pytest import zstandard as zstd @@ -43,15 +42,3 @@ def label(thing, name): def test_compression_wrapper_read(fileobj, compression, filename): wrapped = smart_open.compression.compression_wrapper(fileobj, 'rb', compression, filename) assert wrapped.read() == plain - - -def test_zst_write(): - with tempfile.NamedTemporaryFile(suffix=".zst") as tmp: - with smart_open.open(tmp.name, "wt") as fout: - print("hello world", file=fout) - print("this is a test", file=fout) - - with smart_open.open(tmp.name, "rt") as fin: - got = list(fin) - - assert got == ["hello world\n", "this is a test\n"] diff --git a/smart_open/tests/test_smart_open.py b/smart_open/tests/test_smart_open.py index 067ea19b..4b7c9d32 100644 --- a/smart_open/tests/test_smart_open.py +++ b/smart_open/tests/test_smart_open.py @@ -77,6 +77,18 @@ def named_temporary_file(mode='w+b', prefix=None, suffix=None, delete=True): logger.error(e) +def test_zst_write(): + with tempfile.NamedTemporaryFile(suffix=".zst") as tmp: + with smart_open.open(tmp.name, "wt") as fout: + print("hello world", file=fout) + print("this is a test", file=fout) + + with smart_open.open(tmp.name, "rt") as fin: + got = list(fin) + + assert got == ["hello world\n", "this is a test\n"] + + class ParseUriTest(unittest.TestCase): """ Test ParseUri class. From 27b7f2ed2df9d3eb87d0215f04c3e0cb27fd5aa7 Mon Sep 17 00:00:00 2001 From: Michael Penkov Date: Thu, 21 Mar 2024 16:05:08 +0900 Subject: [PATCH 3/3] use our special named_temporary_file decorator --- smart_open/tests/test_smart_open.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/smart_open/tests/test_smart_open.py b/smart_open/tests/test_smart_open.py index 4b7c9d32..414a9c64 100644 --- a/smart_open/tests/test_smart_open.py +++ b/smart_open/tests/test_smart_open.py @@ -78,7 +78,7 @@ def named_temporary_file(mode='w+b', prefix=None, suffix=None, delete=True): def test_zst_write(): - with tempfile.NamedTemporaryFile(suffix=".zst") as tmp: + with named_temporary_file(suffix=".zst") as tmp: with smart_open.open(tmp.name, "wt") as fout: print("hello world", file=fout) print("this is a test", file=fout)