Added decode_csv and encode_csv methods to io_util.

Author: fabiocaccamo

benedict/utils/io_util.py

@@ -1,8 +1,9 @@
 # -*- coding: utf-8 -*-
 
-from six import binary_type, string_types
+from six import binary_type, string_types, StringIO
 
 import base64
+import csv
 import errno
 import json
 import os
@@ -51,6 +52,28 @@ def decode_base64(s, **kwargs):
     return data
 
 
+def decode_csv(s, **kwargs):
+    # kwargs.setdefault('delimiter', ',')
+    if kwargs.pop('quote', False):
+        kwargs.setdefault('quoting', csv.QUOTE_ALL)
+    columns = kwargs.pop('columns', None)
+    columns_row = kwargs.pop('columns_row', True)
+    f = StringIO(s)
+    r = csv.reader(f, **kwargs)
+    ln = 0
+    data = []
+    for row in r:
+        if ln == 0 and columns_row:
+            if not columns:
+                columns = row
+            ln += 1
+            continue
+        d = dict(zip(columns, row))
+        data.append(d)
+        ln += 1
+    return data
+
+
 def decode_json(s, **kwargs):
     data = json.loads(s, **kwargs)
     return data
@@ -108,6 +131,32 @@ def encode_base64(d, **kwargs):
     return data
 
 
+def encode_csv(l, **kwargs):
+    # kwargs.setdefault('delimiter', ',')
+    if kwargs.pop('quote', False):
+        kwargs.setdefault('quoting', csv.QUOTE_ALL)
+    kwargs.setdefault('lineterminator', '\n')
+    columns = kwargs.pop('columns', None)
+    columns_row = kwargs.pop('columns_row', True)
+    if not columns and len(l) and isinstance(l[0], dict):
+        keys = [str(key) for key in l[0].keys()]
+        columns = list(sorted(keys))
+    f = StringIO()
+    w = csv.writer(f, **kwargs)
+    if columns_row and columns:
+        w.writerow(columns)
+    for item in l:
+        if isinstance(item, dict):
+            row = [item.get(key, '') for key in columns]
+        elif isinstance(item, (list, tuple, set, )):
+            row = item
+        else:
+            row = [item]
+        w.writerow(row)
+    data = f.getvalue()
+    return data
+
+
 def encode_json(d, **kwargs):
     data = json.dumps(d, **kwargs)
     return data