Add tests and fix Python 3 Unicode compatibility.

flask_admin/_compat.py

@@ -31,6 +31,10 @@ if not PY2:
 
         return str(s)
 
+    def csv_encode(s):
+        ''' Returns unicode string expected by Python 3's csv module '''
+        return as_unicode(s)
+
     # Various tools
     from functools import reduce
     from urllib.parse import urljoin, urlparse
@@ -50,6 +54,10 @@ else:
 
         return unicode(s)
 
+    def csv_encode(s):
+        ''' Returns byte string expected by Python 2's csv module '''
+        return as_unicode(s).encode('utf-8')
+
     # Helpers
     reduce = __builtins__['reduce'] if isinstance(__builtins__, dict) else __builtins__.reduce
     from urlparse import urljoin, urlparse

flask_admin/model/base.py

@@ -22,7 +22,7 @@ from flask_admin.helpers import (get_form_data, validate_form_on_submit,
                                  get_redirect_target, flash_errors)
 from flask_admin.tools import rec_getattr
 from flask_admin._backwards import ObsoleteAttr
-from flask_admin._compat import iteritems, OrderedDict, as_unicode
+from flask_admin._compat import iteritems, OrderedDict, as_unicode, csv_encode
 from .helpers import prettify_name, get_mdict_item_or_list
 from .ajax import AjaxModelLoader
 from .fields import ListEditableFieldList
@@ -1962,25 +1962,19 @@ class BaseModelView(BaseView, ActionsMixin):
         writer = csv.writer(Echo())
 
         def generate():
-            # Needed as python 2 csvwriter does not support unicode
-            def fix_unicode(t):
-                return as_unicode(t).encode('utf-8')
-
             # Append the column titles at the beginning
-            titles = [fix_unicode(c[1]) for c in self._list_columns]
+            titles = [csv_encode(c[1]) for c in self._list_columns]
             yield writer.writerow(titles)
 
             for row in data:
-                vals = [fix_unicode(self.get_export_value(row, c[0]))
+                vals = [csv_encode(self.get_export_value(row, c[0]))
                         for c in self._list_columns]
                 yield writer.writerow(vals)
 
-        filename = '{}_{}.csv'.format(
-            self.name,
-            time.strftime("%Y-%m-%d_%H-%M-%S")
-        )
+        filename = '%s_%s.csv' % (self.name,
+                                  time.strftime("%Y-%m-%d_%H-%M-%S"))
 
-        disposition = 'attachment;filename={}'.format(secure_filename(filename))
+        disposition = 'attachment;filename=%s' % (secure_filename(filename),)
 
         return Response(
             stream_with_context(generate()),

flask_admin/tests/test_model.py

@@ -12,6 +12,8 @@ from wtforms import fields
 from flask_admin import Admin, form
 from flask_admin._compat import iteritems, itervalues
 from flask_admin.model import base, filters
+from flask_admin.model.template import macro
+from itertools import islice
 
 
 def wtforms2_and_up(func):
@@ -46,8 +48,8 @@ class SimpleFilter(filters.BaseFilter):
 
 
 class MockModelView(base.BaseModelView):
-    def __init__(self, model, name=None, category=None, endpoint=None, url=None,
-                 **kwargs):
+    def __init__(self, model, data=None, name=None, category=None,
+                 endpoint=None, url=None, **kwargs):
         # Allow to set any attributes from parameters
         for k, v in iteritems(kwargs):
             setattr(self, k, v)
@@ -60,9 +62,12 @@ class MockModelView(base.BaseModelView):
 
         self.search_arguments = []
 
-        self.all_models = {1: Model(1),
-                           2: Model(2)}
-        self.last_id = 3
+        if data is None:
+            self.all_models = {1: Model(1), 2: Model(2)}
+        else:
+            self.all_models = data
+
+        self.last_id = len(self.all_models) + 1
 
     # Scaffolding
     def get_pk_value(self, model):
@@ -89,9 +94,12 @@ class MockModelView(base.BaseModelView):
         return Form
 
     # Data
-    def get_list(self, page, sort_field, sort_desc, search, filters):
+    def get_list(self, page, sort_field, sort_desc, search, filters,
+                 page_size=None):
         self.search_arguments.append((page, sort_field, sort_desc, search, filters))
-        return len(self.all_models), itervalues(self.all_models)
+        count = len(self.all_models)
+        data = islice(itervalues(self.all_models), 0, page_size)
+        return count, data
 
     def get_one(self, id):
         return self.all_models.get(int(id))
@@ -538,3 +546,120 @@ def check_class_name():
 
     view = DummyView(Model)
     eq_(view.name, 'Dummy View')
+
+
+def test_export_csv():
+    app, admin = setup()
+    client = app.test_client()
+
+    # test redirect when csv export is disabled
+    view = MockModelView(Model, column_list=['col1', 'col2'], endpoint="test")
+    admin.add_view(view)
+
+    rv = client.get('/admin/test/export/csv/')
+    eq_(rv.status_code, 302)
+
+    # basic test of csv export with a few records
+    view_data = {
+        1: Model(1, "col1_1", "col2_1"),
+        2: Model(2, "col1_2", "col2_2"),
+        3: Model(3, "col1_3", "col2_3"),
+    }
+
+    view = MockModelView(Model, view_data, can_export=True,
+                         column_list=['col1', 'col2'])
+    admin.add_view(view)
+
+    rv = client.get('/admin/model/export/csv/')
+    data = rv.data.decode('utf-8')
+    eq_(rv.mimetype, 'text/csv')
+    eq_(rv.status_code, 200)
+    ok_("Col1,Col2\r\n"
+        "col1_1,col2_1\r\n"
+        "col1_2,col2_2\r\n"
+        "col1_3,col2_3\r\n" == data)
+
+    # test utf8 characters in csv export
+    view_data[4] = Model(1, u'\u2013ut8_1\u2013', u'\u2013utf8_2\u2013')
+    view = MockModelView(Model, view_data, can_export=True,
+                         column_list=['col1', 'col2'], endpoint="utf8")
+    admin.add_view(view)
+
+    rv = client.get('/admin/utf8/export/csv/')
+    data = rv.data.decode('utf-8')
+    eq_(rv.status_code, 200)
+    ok_(u'\u2013ut8_1\u2013,\u2013utf8_2\u2013\r\n' in data)
+
+    # test row limit
+    view_data = {
+        1: Model(1, "col1_1", "col2_1"),
+        2: Model(2, "col1_2", "col2_2"),
+        3: Model(3, "col1_3", "col2_3"),
+    }
+
+    view = MockModelView(Model, view_data, can_export=True,
+                         column_list=['col1', 'col2'], export_max_rows=2,
+                         endpoint='row_limit_2')
+    admin.add_view(view)
+
+    rv = client.get('/admin/row_limit_2/export/csv/')
+    data = rv.data.decode('utf-8')
+    eq_(rv.status_code, 200)
+    ok_("Col1,Col2\r\n"
+        "col1_1,col2_1\r\n"
+        "col1_2,col2_2\r\n" == data)
+
+    # test None type, integer type, column_labels, and column_formatters
+    view_data = {
+        1: Model(1, "col1_1", 1),
+        2: Model(2, "col1_2", 2),
+        3: Model(3, None, 3),
+    }
+
+    view = MockModelView(
+        Model, view_data, can_export=True, column_list=['col1', 'col2'],
+        column_labels={'col1': 'Str Field', 'col2': 'Int Field'},
+        column_formatters=dict(col2=lambda v, c, m, p: m.col2*2),
+        endpoint="types_and_formatters"
+    )
+    admin.add_view(view)
+
+    rv = client.get('/admin/types_and_formatters/export/csv/')
+    data = rv.data.decode('utf-8')
+    eq_(rv.status_code, 200)
+    ok_("Str Field,Int Field\r\n"
+        "col1_1,2\r\n"
+        "col1_2,4\r\n"
+        ",6\r\n" == data)
+
+    # test column_formatters_export and column_formatters_export
+    type_formatters = {type(None): lambda view, value: "null"}
+
+    view = MockModelView(
+        Model, view_data, can_export=True, column_list=['col1', 'col2'],
+        column_formatters_export=dict(col2=lambda v, c, m, p: m.col2*3),
+        column_formatters=dict(col2=lambda v, c, m, p: m.col2*2),  # overridden
+        column_type_formatters_export=type_formatters,
+        endpoint="export_types_and_formatters"
+    )
+    admin.add_view(view)
+
+    rv = client.get('/admin/export_types_and_formatters/export/csv/')
+    data = rv.data.decode('utf-8')
+    eq_(rv.status_code, 200)
+    ok_("Col1,Col2\r\n"
+        "col1_1,3\r\n"
+        "col1_2,6\r\n"
+        "null,9\r\n" == data)
+
+    # Macros are not implemented for csv export yet and will throw an error
+    view = MockModelView(
+        Model, can_export=True, column_list=['col1', 'col2'],
+        column_formatters=dict(col1=macro('render_macro')),
+        endpoint="macro_exception"
+    )
+    admin.add_view(view)
+
+    rv = client.get('/admin/macro_exception/export/csv/')
+    data = rv.data.decode('utf-8')
+    eq_(rv.status_code, 500)