Implement metadata coding with struct instead of pickle.

This is more portable to implementations in other languages,
if we ever need it.

2 files changed, 51 insertions, 4 deletions

diff --git a/obnamlib/store.py b/obnamlib/store.py
index 8e6226b5..314b4257 100644
--- a/obnamlib/store.py
+++ b/obnamlib/store.py
@@ -23,7 +23,6 @@ import btree
 import errno
 import hashlib
 import os
-import pickle
 import random
 import struct
 import time
@@ -84,12 +83,48 @@ def require_started_generation(method):
     return helper
 
 
+numeric_fields = [x for x in obnamlib.metadata_fields if x.startswith('st_')]
+string_fields = [x for x in obnamlib.metadata_fields 
+                 if x not in numeric_fields]
+num_numeric = len(numeric_fields)
+metadata_format = struct.Struct('!Q' + 'Q' * len(obnamlib.metadata_fields))
+
+
 def encode_metadata(metadata):
-    return pickle.dumps(metadata)
+    flags = 0
+    for i, name in enumerate(obnamlib.metadata_fields):
+        if getattr(metadata, name) is not None:
+            flags |= (1 << i)
+    fields = ([flags] +
+              [getattr(metadata, x) or 0 for x in numeric_fields] +
+              [len(getattr(metadata, x) or '') for x in string_fields])
+    string = ''.join(getattr(metadata, x) or '' for x in string_fields)
+    return metadata_format.pack(*fields) + string
     
 
 def decode_metadata(encoded):
-    return pickle.loads(encoded)
+    buf = buffer(encoded)
+    items = metadata_format.unpack_from(buf)
+
+    flags = items[0]
+    values = list(items[1:len(numeric_fields)+1])
+    lengths = items[len(numeric_fields)+1:]
+    
+    offset = metadata_format.size
+    append = values.append
+    for length in lengths:
+        append(encoded[offset:offset + length])
+        offset += length
+
+    def flagtonone(values):
+        for i, value in enumerate(values):
+            if flags & (1 << i):
+                yield value
+            else:
+                yield None
+
+    args = dict(zip(numeric_fields + string_fields, flagtonone(values)))
+    return obnamlib.Metadata(**args)
 
 
 class NodeStoreVfs(btree.NodeStoreDisk):
diff --git a/obnamlib/store_tests.py b/obnamlib/store_tests.py
index ae6210fa..75f56bfc 100644
--- a/obnamlib/store_tests.py
+++ b/obnamlib/store_tests.py
@@ -45,7 +45,19 @@ class MetadataCodingTests(unittest.TestCase):
         for name in dir(metadata):
             if name in obnamlib.metadata.metadata_fields:
                 self.assertEqual(getattr(metadata, name), 
-                                 getattr(decoded, name))
+                                 getattr(decoded, name),
+                                 'attribute %s must be equal (%s vs %s)' % 
+                                    (name, getattr(metadata, name),
+                                     getattr(decoded, name)))
+
+    def test_round_trip_for_None_values(self):
+        metadata = obnamlib.metadata.Metadata()
+        encoded = obnamlib.store.encode_metadata(metadata)
+        decoded = obnamlib.store.decode_metadata(encoded)
+        for name in dir(metadata):
+            if name in obnamlib.metadata.metadata_fields:
+                self.assertEqual(getattr(decoded, name), None,
+                                 'attribute %s must be None' % name)
 
 
 class ChecksumTreeTests(unittest.TestCase):