Commit 8718c37f authored by Luca Wehrstedt's avatar Luca Wehrstedt Committed by Facebook Github Bot

Support non-byte and non-1d memoryviews when converting to IOBufs

Summary:
Using `shape[0]` as the length of the IOBuf means using the number of _elements_ along the first dimension, rather than the total number of bytes in the memoryview. Thus
- if the element size is not 1 byte, this will get only a fraction of the memory;
- if the memoryview is shapeless, the code will segfault;
- if the memoryview is multidimensional, it will only get a fraction of the memory (and, if the memory is C-contiguous, it will be an odd slice: in a 2x3 memoryview, it will get the entire first row and half of the second).

Perhaps the above was intentional, but I believe that when converting a shaped and typed object to an unshaped and untyped buffer, it's more intuitive to ignore the type and size and just look at the memory backing it.

I also added a check to ensure the view is contiguous. IOBufs can't represent a non-contiguous buffer, except by using a chain, which however in this case could have an arbitrary unbounded length, making it unadvisable. Alternatively we could copy the memory, but this is probably not what the user would expect. Disallowing this seems the safest option.

Reviewed By: yfeldblum

Differential Revision: D18006273

fbshipit-source-id: f21dbb462266eee68f1a5ce0ecea85b1a46e0a8b
parent 328ab26d
......@@ -25,12 +25,14 @@ __all__ = ['IOBuf']
cdef unique_ptr[cIOBuf] from_python_buffer(memoryview view):
"""Take a python object that supports buffer protocol"""
if not view.is_c_contig() and not view.is_f_contig():
raise ValueError("View must be contiguous")
return move(
iobuf_from_python(
get_executor(),
<PyObject*>view,
view.view.buf,
view.shape[0],
view.view.len,
)
)
......
......@@ -13,6 +13,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import array
import struct
import unittest
from folly.iobuf import IOBuf
......@@ -93,3 +95,15 @@ class IOBufTests(unittest.TestCase):
self.assertLessEqual(x, z)
self.assertGreater(y, x)
self.assertGreaterEqual(y, x)
def test_typed(self) -> None:
x = IOBuf(array.array("l", [1, 2, 3, 4, 5])) # type: ignore
self.assertEqual(x.chain_size(), 5 * struct.calcsize("l"))
def test_unshaped(self) -> None:
x = IOBuf(memoryview(b"a").cast("B", shape=[])) # type: ignore
self.assertEqual(x.chain_size(), 1)
def test_multidimensional(self) -> None:
x = IOBuf(memoryview(b"abcdef").cast("B", shape=[3, 2])) # type: ignore
self.assertEqual(x.chain_size(), 6)
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment