Generate Memory Report as valid CSV

rdbtools/__init__.py

@@ -1,6 +1,6 @@
 from rdbtools.parser import RdbCallback, RdbParser, DebugCallback
 from rdbtools.callbacks import JSONCallback, DiffCallback, ProtocolCallback, KeyValsOnlyCallback, KeysOnlyCallback
-from rdbtools.memprofiler import MemoryCallback, PrintAllKeys, StatsAggregator, PrintJustKeys
+from rdbtools.memprofiler import MemoryCallback, PrintAllKeys, MemoryRecord, StatsAggregator, PrintJustKeys
 
 __version__ = '0.1.10'
 VERSION = tuple(map(int, __version__.split('.')))

rdbtools/memprofiler.py

@@ -1,4 +1,5 @@
 import codecs
+import csv,io,sys
 from collections import namedtuple
 import random
 import bisect
@@ -81,10 +82,14 @@ class PrintAllKeys(object):
     def __init__(self, out, bytes, largest):
         self._bytes = bytes
         self._largest = largest
-        self._out = out
-        headers = "%s,%s,%s,%s,%s,%s,%s\n" % (
-            "database", "type", "key", "size_in_bytes", "encoding", "num_elements", "len_largest_element")
-        self._out.write(codecs.encode(headers, 'latin-1'))
+        if sys.version_info < (3,):
+            self._out = out
+        else:
+            #python3 csv expects a text stream
+            self._out = io.TextIOWrapper(out, encoding='utf8',newline='', write_through=True)
+        self._csv = csv.writer(self._out, dialect='excel', lineterminator='\n')
+        self._csv.writerow([ "database", "type", "key", "size_in_bytes",
+            "encoding", "num_elements", "len_largest_element"])
 
         if self._largest is not None:
             self._heap = []
@@ -97,7 +102,9 @@ def next_record(self, record) :
                 rec_str = "%d,%s,%s,%d,%s,%d,%d\n" % (
                     record.database, record.type, record.key, record.bytes, record.encoding, record.size,
                     record.len_largest_element)
-                self._out.write(codecs.encode(rec_str, 'latin-1'))
+                self._csv.writerow([
+                    record.database, record.type, record.key, record.bytes, record.encoding, record.size,
+                    record.len_largest_element])
         else:
             heappush(self._heap, (record.bytes, record))
 

tests/__init__.py

@@ -1,13 +1,14 @@
 import unittest
 from tests.parser_tests import RedisParserTestCase
-from tests.memprofiler_tests import MemoryCallbackTestCase
+from tests.memprofiler_tests import MemoryCallbackTestCase, PrintAllKeysTestCase
 from tests.callbacks_tests import ProtocolTestCase, JsonTestCase, DiffTestCase, KeysTestCase, KeyValsTestCase
 
 
 def all_tests():
     suite = unittest.TestSuite()
     test_case_list = [RedisParserTestCase,
                       MemoryCallbackTestCase,
+                      PrintAllKeysTestCase,
                       ProtocolTestCase,
                       JsonTestCase,
                       DiffTestCase,

tests/memprofiler_tests.py

@@ -1,8 +1,9 @@
 import unittest
 
 from rdbtools import RdbParser
-from rdbtools import MemoryCallback
+from rdbtools import MemoryCallback,PrintAllKeys,MemoryRecord
 import os
+import io
 
 class Stats(object):
     def __init__(self):
@@ -26,3 +27,21 @@ def test_len_largest_element(self):
         stats = get_stats('ziplist_that_compresses_easily.rdb')
 
         self.assertEqual(stats['ziplist_compresses_easily'].len_largest_element, 36, "Length of largest element does not match")
+
+class PrintAllKeysTestCase(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def test_emits_valid_csv(self):
+        stream = io.BytesIO()
+
+        printer = PrintAllKeys(stream, None, None)
+        printer.next_record(MemoryRecord(0, "string", "First,Second", 104, "string", 8, 8))
+        printer.next_record(MemoryRecord(0, "string", 'json:{"key": "value"}', 104, "string", 8, 8))
+
+        expected_csv = b"""database,type,key,size_in_bytes,encoding,num_elements,len_largest_element
+0,string,"First,Second",104,string,8,8
+0,string,"json:{""key"": ""value""}",104,string,8,8
+"""
+        self.assertEqual(stream.getvalue(), expected_csv)
+