ENH: reimplement _move_into_mutable_buffer in C.

By writing our move function in C we can hide the original bytes object
from the user while still ensuring that the lifetime is managed
correctly. This implementation is designed to make it impossible to get
access to the invalid bytes object from pure python.

Author: Joe Jevnik

pandas/io/packers.py

@@ -41,7 +41,7 @@
 from datetime import datetime, date, timedelta
 from dateutil.parser import parse
 import os
-import sys
+from textwrap import dedent
 import warnings
 
 import numpy as np
@@ -64,6 +64,51 @@
 import pandas.core.internals as internals
 
 from pandas.msgpack import Unpacker as _Unpacker, Packer as _Packer, ExtType
+from pandas.util._move import (
+    BadMove as _BadMove,
+    move_into_mutable_buffer as _move_into_mutable_buffer,
+)
+
+# check whcih compression libs we have installed
+try:
+    import zlib
+
+    def _check_zlib():
+        pass
+except ImportError:
+    def _check_zlib():
+        raise ValueError('zlib is not installed')
+
+_check_zlib.__doc__ = dedent(
+    """\
+    Check if zlib is installed.
+
+    Raises
+    ------
+    ValueError
+        Raised when zlib is not installed.
+    """,
+)
+
+try:
+    import blosc
+
+    def _check_blosc():
+        pass
+except ImportError:
+    def _check_blosc():
+        raise ValueError('blosc is not installed')
+
+_check_blosc.__doc__ = dedent(
+    """\
+    Check if blosc is installed.
+
+    Raises
+    ------
+    ValueError
+        Raised when blosc is not installed.
+    """,
+)
 
 # until we can pass this into our conversion functions,
 # this is pretty hacky
@@ -221,102 +266,31 @@ def convert(values):
         return v.tolist()
 
     if compressor == 'zlib':
+        _check_zlib()
 
         # return string arrays like they are
         if dtype == np.object_:
             return v.tolist()
 
         # convert to a bytes array
         v = v.tostring()
-        import zlib
         return ExtType(0, zlib.compress(v))
 
     elif compressor == 'blosc':
+        _check_blosc()
 
         # return string arrays like they are
         if dtype == np.object_:
             return v.tolist()
 
         # convert to a bytes array
         v = v.tostring()
-        import blosc
         return ExtType(0, blosc.compress(v, typesize=dtype.itemsize))
 
     # ndarray (on original dtype)
     return ExtType(0, v.tostring())
 
 
-class _BadMove(ValueError):
-    """Exception used to indicate that a move was attempted on a value with
-    more than a single reference.
-
-    Parameters
-    ----------
-    data : any
-        The data which was passed to ``_move_into_mutable_buffer``.
-
-    See Also
-    --------
-    _move_into_mutable_buffer
-    """
-    def __init__(self, data):
-        self.data = data
-
-    def __str__(self):
-        return 'cannot move data from a named object'
-
-
-def _move_into_mutable_buffer(bytes_rvalue):
-    """Moves a bytes object that is about to be destroyed into a mutable buffer
-    without copying the data.
-
-    Parameters
-    ----------
-    bytes_rvalue : bytes with 1 refcount.
-        The bytes object that you want to move into a mutable buffer. This
-        cannot be a named object. It must only have a single reference.
-
-    Returns
-    -------
-    buf : memoryview
-        A mutable buffer that was previously used as that data for
-        ``bytes_rvalue``.
-
-    Raises
-    ------
-    _BadMove
-        Raised when a move is attempted on an object with more than one
-        reference.
-
-    Notes
-    -----
-    If you want to use this function you are probably wrong.
-    """
-    if sys.getrefcount(bytes_rvalue) != 3:
-        # The three references are:
-        # 1. The callers stack (this is the only external reference)
-        # 2. The locals for this function
-        # 3. This function's stack (to pass to `sys.getrefcount`)
-        raise _BadMove(bytes_rvalue)
-
-    # create a numpy array from the memory of `bytes_rvalue`
-    arr = np.frombuffer(bytes_rvalue, dtype=np.int8)
-    try:
-        # mark this array as mutable
-        arr.flags.writeable = True
-        # At this point any mutations to `arr` will invalidate `bytes_rvalue`.
-        # This would be fine but `np.frombuffer` is going to store this object
-        # on `arr.base`. In order to preserve user's sanity we are going to
-        # destroy `arr` to drop the final reference to `bytes_rvalue` and just
-        # return a `memoryview` of the now mutable data. This dance is very
-        # fast and makes it impossible for users to shoot themselves in the
-        # foot.
-        return memoryview(arr)
-    finally:
-        # assure that our mutable view is destroyed even if we raise
-        del arr
-
-
 def unconvert(values, dtype, compress=None):
 
     as_is_ext = isinstance(values, ExtType) and values.code == 0
@@ -334,9 +308,11 @@ def unconvert(values, dtype, compress=None):
 
     if compress:
         if compress == u'zlib':
-            from zlib import decompress
+            _check_zlib()
+            decompress = zlib.decompress
         elif compress == u'blosc':
-            from blosc import decompress
+            _check_blosc()
+            decompress = blosc.decompress
         else:
             raise ValueError("compress must be one of 'zlib' or 'blosc'")
 
@@ -350,7 +326,7 @@ def unconvert(values, dtype, compress=None):
             # We don't just store this in the locals because we want to
             # minimize the risk of giving users access to a `bytes` object
             # whose data is also given to a mutable buffer.
-            values = e.data
+            values = e.args[0]
             if len(values) > 1:
                 # The empty string and single characters are memoized in many
                 # string creating functions in the capi. This case should not

pandas/tests/test_util.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 import nose
 
+from pandas.util._move import move_into_mutable_buffer, BadMove
 from pandas.util.decorators import deprecate_kwarg
 from pandas.util.validators import validate_args, validate_kwargs
 
@@ -150,6 +151,38 @@ def test_validation(self):
         kwargs = {'f': 'foo', 'b': 'bar'}
         validate_kwargs('func', kwargs, *compat_args)
 
+
+class TestMove(tm.TestCase):
+    def test_more_than_one_ref(self):
+        """Test case for when we try to use ``move_into_mutable_buffer`` when
+        the object being moved has other references.
+        """
+        b = b'testing'
+
+        with tm.assertRaises(BadMove) as e:
+            def handle_success(type_, value, tb):
+                self.assertIs(value.args[0], b)
+                return type(e).handle_success(e, type_, value, tb)  # super
+
+            e.handle_success = handle_success
+            move_into_mutable_buffer(b)
+
+    def test_exactly_one_ref(self):
+        """Test case for when the object being moved has exactly one reference.
+        """
+        b = b'testing'
+
+        # We need to pass an expression on the stack to ensure that there are
+        # not extra references hanging around. We cannot rewrite this test as
+        #   buf = b[:-3]
+        #   as_stolen_buf = move_into_mutable_buffer(buf)
+        # because then we would have more than one reference to buf.
+        as_stolen_buf = move_into_mutable_buffer(b[:-3])
+
+        # materialize as bytearray to show that it is mutable
+        self.assertEqual(bytearray(as_stolen_buf), b'test')
+
+
 if __name__ == '__main__':
     nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
                    exit=False)