tweak binary funcs param checking + tests

Author: rossengeorgiev

tests/test_binary_vdf.py

@@ -65,6 +65,16 @@ def test_dumps_unicode(self):
     def test_dumps_unicode_alternative(self):
         self.assertEqual(vdf.binary_dumps({u('a'): u('b')}, alt_format=True), b'\x01a\x00b\x00\x0b')
 
+    def test_dump_params_invalid(self):
+        with self.assertRaises(TypeError):
+            vdf.binary_dump([], BytesIO())
+            vdf.binary_dump({}, b'aaaa')
+
+    def test_dumps_params_invalid(self):
+        with self.assertRaises(TypeError):
+            vdf.binary_dumps([])
+            vdf.binary_dumps(b'aaaa')
+
     def test_dumps_key_invalid_type(self):
         with self.assertRaises(TypeError):
             vdf.binary_dumps({1:1})
@@ -81,6 +91,19 @@ def test_alternative_format(self):
         with self.assertRaises(SyntaxError):
             vdf.binary_loads(b'\x00a\x00\x00b\x00\x08\x08', alt_format=True)
 
+    def test_load_params_invalid(self):
+        with self.assertRaises(TypeError):
+            vdf.binary_load(b'aaaa')
+            vdf.binary_load(1234)
+            vdf.binary_load(BytesIO(b'aaaa'), b'bbbb')
+
+    def test_loads_params_invalid(self):
+        with self.assertRaises(TypeError):
+            vdf.binary_loads([])
+            vdf.binary_loads(11111)
+            vdf.binary_loads(BytesIO())
+            vdf.binary_load(b'', b'bbbb')
+
     def test_loads_unbalanced_nesting(self):
         with self.assertRaises(SyntaxError):
             vdf.binary_loads(b'\x00a\x00\x00b\x00\x08')

tests/test_vdf.py

@@ -84,6 +84,14 @@ def setUp(self):
     def tearDown(self):
         self.f.close()
 
+    def test_dump_params_invalid(self):
+        with self.assertRaises(TypeError):
+            # pretty/escaped only accept bool
+            vdf.dump({'a': 1}, StringIO(), pretty=1)
+            vdf.dumps({'a': 1}, pretty=1)
+            vdf.dump({'a': 1}, StringIO(), escaped=1)
+            vdf.dumps({'a': 1}, escaped=1)
+
     def test_routine_dumps_asserts(self):
         for x in [5, 5.5, 1.0j, True, None, (), {}, lambda: 0, sys.stdin, self.f]:
             for y in [5, 5.5, 1.0j, None, [], (), {}, lambda: 0, sys.stdin, self.f]:

vdf/__init__.py

@@ -8,8 +8,14 @@
 import sys
 import struct
 from binascii import crc32
-from io import StringIO as unicodeIO
 from io import BytesIO
+from io import StringIO as unicodeIO
+
+try:
+    from collections.abc import Mapping
+except:
+    from collections import Mapping
+
 from vdf.vdict import VDFDict
 
 # Py2 & Py3 compatibility
@@ -72,7 +78,7 @@ def parse(fp, mapper=dict, merge_duplicate_keys=True, escaped=True):
     same key into one instead of overwriting. You can se this to ``False`` if you are
     using ``VDFDict`` and need to preserve the duplicates.
     """
-    if not issubclass(mapper, dict):
+    if not issubclass(mapper, Mapping):
         raise TypeError("Expected mapper to be subclass of dict, got %s" % type(mapper))
     if not hasattr(fp, 'readline'):
         raise TypeError("Expected fp to be a file-like object supporting line iteration")
@@ -200,7 +206,7 @@ def dumps(obj, pretty=False, escaped=True):
     """
     Serialize ``obj`` to a VDF formatted ``str``.
     """
-    if not isinstance(obj, dict):
+    if not isinstance(obj, Mapping):
         raise TypeError("Expected data to be an instance of``dict``")
     if not isinstance(pretty, bool):
         raise TypeError("Expected pretty to be of type bool")
@@ -215,7 +221,7 @@ def dump(obj, fp, pretty=False, escaped=True):
     Serialize ``obj`` as a VDF formatted stream to ``fp`` (a
     ``.write()``-supporting file-like object).
     """
-    if not isinstance(obj, dict):
+    if not isinstance(obj, Mapping):
         raise TypeError("Expected data to be an instance of``dict``")
     if not hasattr(fp, 'write'):
         raise TypeError("Expected fp to have write() method")
@@ -239,7 +245,7 @@ def _dump_gen(data, pretty=False, escaped=True, level=0):
         if escaped and isinstance(key, string_type):
             key = _escape(key)
 
-        if isinstance(value, dict):
+        if isinstance(value, Mapping):
             yield '%s"%s"\n%s{\n' % (line_indent, key, line_indent)
             for chunk in _dump_gen(value, pretty, escaped, level+1):
                 yield chunk
@@ -295,12 +301,8 @@ def binary_loads(b, mapper=dict, merge_duplicate_keys=True, alt_format=False, ra
     """
     if not isinstance(b, bytes):
         raise TypeError("Expected s to be bytes, got %s" % type(b))
-    if not issubclass(mapper, dict):
-        raise TypeError("Expected mapper to be subclass of dict, got %s" % type(mapper))
-
-    fp = BytesIO(b)
 
-    return binary_load(fp, mapper, merge_duplicate_keys, alt_format, raise_on_remaining)
+    return binary_load(BytesIO(b), mapper, merge_duplicate_keys, alt_format, raise_on_remaining)
 
 def binary_load(fp, mapper=dict, merge_duplicate_keys=True, alt_format=False, raise_on_remaining=False):
     """
@@ -317,7 +319,7 @@ def binary_load(fp, mapper=dict, merge_duplicate_keys=True, alt_format=False, ra
     """
     if not hasattr(fp, 'read') or not hasattr(fp, 'tell') or not hasattr(fp, 'seek'):
         raise TypeError("Expected fp to be a file-like object with tell()/seek() and read() returning bytes")
-    if not issubclass(mapper, dict):
+    if not issubclass(mapper, Mapping):
         raise TypeError("Expected mapper to be subclass of dict, got %s" % type(mapper))
 
     # helpers
@@ -415,12 +417,16 @@ def binary_dumps(obj, alt_format=False):
     """
     Serialize ``obj`` to a binary VDF formatted ``bytes``.
     """
-    return b''.join(_binary_dump_gen(obj, alt_format=alt_format))
+    buf = BytesIO()
+    binary_dump(obj, buf, alt_format)
+    return buf.getvalue()
 
 def binary_dump(obj, fp, alt_format=False):
     """
     Serialize ``obj`` to a binary VDF formatted ``bytes`` and write it to ``fp`` filelike object
     """
+    if not isinstance(obj, Mapping):
+        raise TypeError("Expected obj to be type of Mapping")
     if not hasattr(fp, 'write'):
         raise TypeError("Expected fp to have write() method")
 
@@ -442,7 +448,7 @@ def _binary_dump_gen(obj, level=0, alt_format=False):
         else:
             raise TypeError("dict keys must be of type str, got %s" % type(key))
 
-        if isinstance(value, dict):
+        if isinstance(value, Mapping):
             yield BIN_NONE + key + BIN_NONE
             for chunk in _binary_dump_gen(value, level+1, alt_format=alt_format):
                 yield chunk