Repository: piskvorky/smart_open

Issue:
zstd write does not work with `wb` mode
#### Problem description
FileLikeProxy, and therefore smart_open, produces an AttributeError when trying to open a zstd S3 file in `wb` mode

Cousin to https://github.com/piskvorky/smart_open/issues/810 maybe?

#### Steps/code to reproduce the problem

```
>>> smart_open.open('s3://my-real-bucket-was-here/a-path.zst', mode='wb')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/root/.pyenv/versions/3.11.6/lib/python3.11/site-packages/smart_open/smart_open_lib.py", line 251, in open
    return so_utils.FileLikeProxy(decoded, binary)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.pyenv/versions/3.11.6/lib/python3.11/site-packages/smart_open/utils.py", line 213, in __init__
    self.__inner = inner
    ^^^^^^^^^^^^
AttributeError: 'zstd.ZstdCompressionWriter' object has no attribute '_FileLikeProxy__inner'
```

Only occurs with `wb` mode, not `w` mode

#### Versions

```python
>>> import platform, sys, smart_open
>>> print(platform.platform())
Linux-5.4.209-116.367.amzn2.x86_64-x86_64-with-glibc2.35
>>> print("Python", sys.version)
Python 3.11.6 (main, Mar 11 2024, 06:41:25) [GCC 11.4.0]
>>> print("smart_open", smart_open.__version__)
smart_open 7.0.3
>>> print("wrapt", wrapt.__version__)
wrapt 1.15.0
```

#### Checklist

Before you create the issue, please make sure you have:

- [x] Described the problem clearly
- [x] Provided a minimal reproducible example, including any required data
- [x] Provided the version numbers of the relevant software



Discussion:

Patch:
diff --git a/smart_open/compression.py b/smart_open/compression.py
index 1c88224e..e7c6d93c 100644
--- a/smart_open/compression.py
+++ b/smart_open/compression.py
@@ -6,6 +6,7 @@
 # from the MIT License (MIT).
 #
 """Implements the compression layer of the ``smart_open`` library."""
+import io
 import logging
 import os.path
 
@@ -108,6 +109,17 @@ def _handle_gzip(file_obj, mode):
 def _handle_zstd(file_obj, mode):
     import zstandard  # type: ignore
     result = zstandard.open(filename=file_obj, mode=mode)
+    # zstandard.open returns an io.TextIOWrapper in text mode, but otherwise
+    # returns a raw stream reader/writer, and we need the `io` wrapper
+    # to make FileLikeProxy work correctly.
+    #
+    # See:
+    #
+    # https://github.com/indygreg/python-zstandard/blob/d7d81e79dbe74feb22fb73405ebfb3e20f4c4653/zstandard/__init__.py#L169-L174
+    if "b" in mode and "w" in mode:
+        result = io.BufferedWriter(result)
+    elif "b" in mode and "r" in mode:
+        result = io.BufferedReader(result)
     return result
 
 


Test patch:
diff --git a/smart_open/tests/test_smart_open.py b/smart_open/tests/test_smart_open.py
index 414a9c64..e3727271 100644
--- a/smart_open/tests/test_smart_open.py
+++ b/smart_open/tests/test_smart_open.py
@@ -89,6 +89,18 @@ def test_zst_write():
     assert got == ["hello world\n", "this is a test\n"]
 
 
+def test_zst_write_binary():
+    with named_temporary_file(suffix=".zst") as tmp:
+        with smart_open.open(tmp.name, "wb") as fout:
+            fout.write(b"hello world\n")
+            fout.write(b"this is a test\n")
+
+        with smart_open.open(tmp.name, "rb") as fin:
+            got = list(fin)
+
+    assert got == [b"hello world\n", b"this is a test\n"]
+
+
 class ParseUriTest(unittest.TestCase):
     """
     Test ParseUri class.