Merge pull request #1557 from Scifabric/improve-csv-exporter

Improve csv exporter

Author: teleyinex

pybossa/exporter/csv_export.py

@@ -28,96 +28,28 @@
 from pybossa.util import UnicodeWriter
 from werkzeug.datastructures import FileStorage
 from werkzeug.utils import secure_filename
+from flatten_json import flatten
+import pandas as pd
 
 
 class CsvExporter(Exporter):
 
-    def _format_csv_row(self, row, ty):
-        tmp = row.keys()
-        task_keys = []
-        for k in tmp:
-            k = "%s__%s" % (ty, k)
-            task_keys.append(k)
-        if (type(row['info']) == dict):
-            task_info_keys = []
-            tmp = row['info'].keys()
-            for k in tmp:
-                k = "%sinfo__%s" % (ty, k)
-                task_info_keys.append(k)
-        else:
-            task_info_keys = []
-
-        keys = sorted(task_keys + task_info_keys)
-        values = []
-        _prefix = "%sinfo" % ty
-        for k in keys:
-            prefix, k = k.split("__")
-            if prefix == _prefix:
-                if row['info'].get(k) is not None:
-                    values.append(row['info'][k])
-                else:
-                    values.append(None)
-            else:
-                if row.get(k) is not None:
-                    values.append(row[k])
-                else:
-                    values.append(None)
-
-        return values
-
-    def _handle_row(self, writer, t, ty):
-        normal_ty = filter(lambda char: char.isalpha(), ty)
-        writer.writerow(self._format_csv_row(t.dictize(), ty=normal_ty))
-
-    def _get_csv(self, out, writer, table, id):
-        for tr in getattr(task_repo, 'filter_%ss_by' % table)(project_id=id,
-                                                              yielded=True):
-            self._handle_row(writer, tr, table)
-        out.seek(0)
-        yield out.read()
-
-    def _format_headers(self, t, ty):
-        tmp = t.dictize().keys()
-        task_keys = []
-        for k in tmp:
-            k = "%s__%s" % (ty, k)
-            task_keys.append(k)
-        if (type(t.info) == dict):
-            task_info_keys = []
-            tmp = t.info.keys()
-            for k in tmp:
-                k = "%sinfo__%s" % (ty, k)
-                task_info_keys.append(k)
-        else:
-            task_info_keys = []
-        keys = task_keys + task_info_keys
-        return sorted(keys)
-
     def _respond_csv(self, ty, id):
         out = tempfile.TemporaryFile()
         writer = UnicodeWriter(out)
-        t = getattr(task_repo, 'get_%s_by' % ty)(project_id=id)
-        if t is not None:
-            headers = self._format_headers(t, ty)
-            writer.writerow(headers)
-
-            return self._get_csv(out, writer, ty, id)
-        else:
-            def empty_csv(out):
-                yield out.read()
-            return empty_csv(out)
-
+        data = getattr(task_repo, 'filter_%ss_by' % ty)(project_id=id)
+        flat_data = [flatten(datum.dictize()) for datum in data]
+        return pd.DataFrame(flat_data)
+        
     def _make_zip(self, project, ty):
         name = self._project_name_latin_encoded(project)
-        csv_task_generator = self._respond_csv(ty, project.id)
-        if csv_task_generator is not None:
-            # TODO: use temp file from csv generation directly
+        dataframe = self._respond_csv(ty, project.id)
+        if dataframe is not None:
             datafile = tempfile.NamedTemporaryFile()
             try:
-                for line in csv_task_generator:
-                    datafile.write(str(line))
+                dataframe.to_csv(datafile, index=False,
+                                 encoding='utf-8')
                 datafile.flush()
-                csv_task_generator.close()  # delete temp csv file
                 zipped_datafile = tempfile.NamedTemporaryFile()
                 try:
                     _zip = self._zip_factory(zipped_datafile.name)

setup.py

@@ -58,11 +58,13 @@
     "webassets>=0.12.1, <0.12.2",
     "readability-lxml>=0.6.2, <1.0",
     "pybossa-onesignal",
+    "pandas>=0.20.2, <0.20.3",
+    "flatten-json>=0.1.5, <0.1.6"
 ]
 
 setup(
     name = 'pybossa',
-    version = '2.4.2',
+    version = '2.4.3',
     packages = find_packages(),
     install_requires = requirements,
     # only needed when installing directly from setup.py (PyPi, eggs?) and pointing to e.g. a git repo.

test/test_web.py

@@ -43,6 +43,7 @@
 from unidecode import unidecode
 from werkzeug.utils import secure_filename
 from nose.tools import assert_raises
+from flatten_json import flatten
 
 
 class TestWeb(web.Helper):
@@ -4267,29 +4268,30 @@ def test_export_task_csv(self):
         assert len(exported_tasks) == len(project.tasks), err_msg
         for t in project.tasks:
             err_msg = "All the task column names should be included"
-            for tk in t.dictize().keys():
-                expected_key = "task__%s" % tk
+            for tk in flatten(t.dictize()).keys():
+                expected_key = "%s" % tk
                 assert expected_key in keys, err_msg
             err_msg = "All the task.info column names should be included"
             for tk in t.info.keys():
-                expected_key = "taskinfo__%s" % tk
+                expected_key = "info_%s" % tk
                 assert expected_key in keys, err_msg
 
         for et in exported_tasks:
-            task_id = et[keys.index('task__id')]
+            task_id = et[keys.index('id')]
             task = db.session.query(Task).get(task_id)
+            task_dict_flat = flatten(task.dictize())
             task_dict = task.dictize()
-            for k in task_dict:
-                slug = 'task__%s' % k
-                err_msg = "%s != %s" % (task_dict[k], et[keys.index(slug)])
-                if k != 'info':
-                    assert unicode(task_dict[k]) == et[keys.index(slug)], err_msg
+            for k in task_dict_flat.keys():
+                slug = '%s' % k
+                err_msg = "%s != %s" % (task_dict_flat[k], et[keys.index(slug)])
+                if task_dict_flat[k] is not None:
+                    assert unicode(task_dict_flat[k]) == et[keys.index(slug)], err_msg
                 else:
-                    assert json.dumps(task_dict[k]) == et[keys.index(slug)], err_msg
+                    assert u'' == et[keys.index(slug)], err_msg
             for k in task_dict['info'].keys():
-                slug = 'taskinfo__%s' % k
+                slug = 'info_%s' % k
                 err_msg = "%s != %s" % (task_dict['info'][k], et[keys.index(slug)])
-                assert unicode(task_dict['info'][k]) == et[keys.index(slug)], err_msg
+                assert unicode(task_dict_flat[slug]) == et[keys.index(slug)], err_msg
         # Tasks are exported as an attached file
         content_disposition = 'attachment; filename=%d_project1_task_csv.zip' % project.id
         assert res.headers.get('Content-Disposition') == content_disposition, res.headers
@@ -4363,28 +4365,21 @@ def test_53_export_task_runs_csv(self):
         assert len(exported_task_runs) == len(project.task_runs), err_msg
 
         for t in project.tasks[0].task_runs:
-            for tk in t.dictize().keys():
-                expected_key = "task_run__%s" % tk
-                assert expected_key in keys, expected_key
-            for tk in t.info.keys():
-                expected_key = "task_runinfo__%s" % tk
+            for tk in flatten(t.dictize()).keys():
+                expected_key = "%s" % tk
                 assert expected_key in keys, expected_key
 
         for et in exported_task_runs:
-            task_run_id = et[keys.index('task_run__id')]
+            task_run_id = et[keys.index('id')]
             task_run = db.session.query(TaskRun).get(task_run_id)
-            task_run_dict = task_run.dictize()
+            task_run_dict = flatten(task_run.dictize())
             for k in task_run_dict:
-                slug = 'task_run__%s' % k
+                slug = '%s' % k
                 err_msg = "%s != %s" % (task_run_dict[k], et[keys.index(slug)])
-                if k != 'info':
+                if task_run_dict[k] is not None:
                     assert unicode(task_run_dict[k]) == et[keys.index(slug)], err_msg
                 else:
-                    assert json.dumps(task_run_dict[k]) == et[keys.index(slug)], err_msg
-            for k in task_run_dict['info'].keys():
-                slug = 'task_runinfo__%s' % k
-                err_msg = "%s != %s" % (task_run_dict['info'][k], et[keys.index(slug)])
-                assert unicode(task_run_dict['info'][k]) == et[keys.index(slug)], err_msg
+                    assert u'' == et[keys.index(slug)], err_msg
         # Task runs are exported as an attached file
         content_disposition = 'attachment; filename=%d_project1_task_run_csv.zip' % project.id
         assert res.headers.get('Content-Disposition') == content_disposition, res.headers