Return number of bytes written in case of BlockingIOError.
In case of a partial write BufferedIOBase doesn't return the length of data written, but gives a BlockingIOError that contains the number of characters written.
See https://docs.python.org/3/library/io.html#io.BufferedIOBase.write
Also, RawIOBase returns None (and not 0) when no data can be written.
PiperOrigin-RevId: 268167212
diff --git a/python/util/file_object_adapter.py b/python/util/file_object_adapter.py
index 727c4b0..42f1cd4 100644
--- a/python/util/file_object_adapter.py
+++ b/python/util/file_object_adapter.py
@@ -20,6 +20,8 @@
from __future__ import google_type_annotations
from __future__ import print_function
+import io
+
from tink.python.cc.clif import simple_output_stream
@@ -35,7 +37,11 @@
def write(self, data: bytes) -> int:
"""Writes to underlying file object and returns number of bytes written."""
- return self._file_object.write(data)
+ try:
+ written = self._file_object.write(data)
+ return 0 if written is None else written
+ except io.BlockingIOError as e:
+ return e.characters_written
def close(self) -> None:
self._file_object.close()
diff --git a/python/util/file_object_adapter_test.py b/python/util/file_object_adapter_test.py
index d83d9f3..e49a053 100644
--- a/python/util/file_object_adapter_test.py
+++ b/python/util/file_object_adapter_test.py
@@ -18,6 +18,7 @@
import io
from absl.testing import absltest
+import mock
from tink.python.util import file_object_adapter
@@ -57,41 +58,35 @@
adapter.close()
def test_non_writable(self):
+ file_object = mock.Mock()
+ file_object.writable = mock.Mock(return_value=False)
- class TestNonWritableObject(io.RawIOBase):
- """Test non-writable file-like object."""
-
- def writable(self):
- return False
-
- non_writable_object = TestNonWritableObject()
self.assertRaises(TypeError, file_object_adapter.FileObjectAdapter,
- non_writable_object)
+ file_object)
- def test_partial_write(self):
+ def test_write_returns_none(self):
+ file_object = mock.Mock()
+ file_object.writable = mock.Mock(return_value=True)
+ file_object.write = mock.Mock(return_value=None)
- class TestFileObject(io.RawIOBase):
- """Test file-like object that always writes only first 5 bytes of data."""
+ adapter = file_object_adapter.FileObjectAdapter(file_object)
+ self.assertEqual(0, adapter.write(b'something'))
- def __init__(self):
- super(TestFileObject, self).__init__()
- self.value = b''
+ def test_write_raises_blocking_error(self):
+ file_object = mock.Mock()
+ file_object.writable = mock.Mock(return_value=True)
+ file_object.write = mock.Mock(side_effect=io.BlockingIOError(None, None, 5))
- def writable(self):
- return True
-
- def write(self, data):
- self.value += data[:5]
- return 5
-
- def tell(self):
- return len(self.value)
-
- file_object = TestFileObject()
adapter = file_object_adapter.FileObjectAdapter(file_object)
self.assertEqual(5, adapter.write(b'something'))
- self.assertEqual(5, adapter.position())
- self.assertEqual(b'somet', file_object.value)
+
+ def test_partial_write(self):
+ file_object = mock.Mock()
+ file_object.writable = mock.Mock(return_value=True)
+ file_object.write = mock.Mock(wraps=lambda data: len(data) - 1)
+
+ adapter = file_object_adapter.FileObjectAdapter(file_object)
+ self.assertEqual(8, adapter.write(b'something'))
if __name__ == '__main__':
