Merge pull request #101 from xrotwang/master

Ported freezing as csv to python 3.4, thanks @xrotwang

dataset/freeze/config.py

@@ -1,10 +1,7 @@
 import json
 import yaml
 
-try:
-    str = unicode
-except NameError:
-    pass
+from six import text_type, PY3
 
 from dataset.util import FreezeException
 
@@ -33,14 +30,17 @@ class Configuration(object):
         extension = file_name.rsplit('.', 1)[-1]
         loader = DECODER.get(extension, json)
         try:
-            fh = open(file_name, 'rb')
+            if loader == json and PY3:  # pragma: no cover
+                fh = open(file_name, encoding='utf8')
+            else:
+                fh = open(file_name, 'rb')
             try:
                 self.data = loader.load(fh)
             except ValueError as ve:
                 raise FreezeException("Invalid freeze file: %s" % ve)
             fh.close()
         except IOError as ioe:
-            raise FreezeException(str(ioe))
+            raise FreezeException(text_type(ioe))
 
     @property
     def exports(self):
@@ -64,7 +64,7 @@ class Export(object):
     def get_normalized(self, name, default=None):
         value = self.get(name, default=default)
         if value not in [None, default]:
-            value = str(value).lower().strip()
+            value = text_type(value).lower().strip()
         return value
 
     def get_bool(self, name, default=False):

dataset/freeze/format/common.py

@@ -3,10 +3,7 @@ import re
 import sys
 import locale
 
-try:
-    str = unicode
-except NameError:
-    pass
+from six import binary_type, text_type
 
 from dataset.util import FreezeException
 from slugify import slugify
@@ -16,7 +13,7 @@ TMPL_KEY = re.compile("{{([^}]*)}}")
 
 OPERATIONS = {
         'identity': lambda x: x,
-        'lower': lambda x: str(x).lower(),
+        'lower': lambda x: text_type(x).lower(),
         'slug': slugify
         }
 
@@ -24,6 +21,7 @@ OPERATIONS = {
 class Serializer(object):
 
     def __init__(self, export, query):
+        self._encoding = locale.getpreferredencoding()
         self.export = export
         self.query = query
         self._paths = []
@@ -35,10 +33,14 @@ class Serializer(object):
 
     def _get_basepath(self):
         prefix = self.export.get('prefix', '')
+        if isinstance(prefix, binary_type):
+            prefix = text_type(prefix, encoding=self._encoding)
         prefix = os.path.abspath(prefix)
         prefix = os.path.realpath(prefix)
         self._prefix = prefix
         filename = self.export.get('filename')
+        if isinstance(filename, binary_type):
+            filename = text_type(filename, encoding=self._encoding)
         if filename is None:
             raise FreezeException("No 'filename' is specified")
         self._basepath = os.path.join(prefix, filename)
@@ -50,8 +52,7 @@ class Serializer(object):
                 op, key = key.split(':', 1)
             return str(OPERATIONS.get(op)(data.get(key, '')))
         path = TMPL_KEY.sub(repl, self._basepath)
-        enc = locale.getpreferredencoding()
-        return os.path.realpath(path.encode(enc, 'replace'))
+        return os.path.realpath(path)
 
     def file_name(self, row):
         # signal that there is a fileobj available:

dataset/freeze/format/fcsv.py

@@ -1,13 +1,22 @@
+from __future__ import unicode_literals
 import csv
-from datetime import datetime
+from datetime import datetime, date
+
+from six import PY3, text_type
 
 from dataset.freeze.format.common import Serializer
 
 
 def value_to_str(value):
-    if isinstance(value, datetime):
-        return value.isoformat()
-    if hasattr(value, 'encode'):
+    if isinstance(value, (datetime, date)):
+        #
+        # FIXME: this check does not work for values returned from a db query!
+        # As a workaround, we make sure, the isoformat call returns the regular
+        # str representation.
+        #
+        sep = ' ' if PY3 else str(' ')
+        return text_type(value.isoformat(sep=sep))
+    if not PY3 and hasattr(value, 'encode'):
         return value.encode('utf-8')
     if value is None:
         return ''
@@ -25,12 +34,18 @@ class CSVSerializer(Serializer):
 
             # handle fileobj that has been passed in:
             if path is not None:
-                fh = open(path, 'wb')
+                if PY3:
+                    fh = open(path, 'wt', encoding='utf8', newline='')
+                else:
+                    fh = open(path, 'wb')
             else:
                 fh = self.fileobj
 
             writer = csv.writer(fh)
-            writer.writerow([k.encode('utf-8') for k in keys])
+            if PY3:
+                writer.writerow(keys)
+            else:
+                writer.writerow([k.encode('utf-8') for k in keys])
             self.handles[path] = (writer, fh)
         writer, fh = self.handles[path]
         values = [value_to_str(result.get(k)) for k in keys]

setup.py

@@ -30,6 +30,7 @@ setup(
     include_package_data=False,
     zip_safe=False,
     install_requires=[
+        'six',
         'sqlalchemy >= 0.9.1',
         'alembic >= 0.6.2',
         'python-slugify >= 0.0.6',

test/Freezefile.yaml

@@ -0,0 +1,32 @@
+common:
+
+  database: "postgresql://user:password@localhost/operational_database"
+  prefix: my_project/dumps/
+  format: json
+  nested:
+
+    property: "inner"
+
+exports:
+
+  - query: "SELECT id, title, date FROM events"
+    filename: "index.json"
+    number: 5
+    bool: true
+    nested:
+
+      property: "override"
+
+  - query: "SELECT id, title, date, country FROM events"
+    filename: "countries/{{country}}.csv"
+    format: csv
+
+  - query: "SELECT * FROM events"
+    filename: "events/{{id}}.json"
+    mode: item
+    wrap: true
+
+  - query: "SELECT * FROM events"
+    filename: "all.json"
+    format: tabson
+

test/test_freeze.py

@@ -0,0 +1,62 @@
+# coding: utf8
+from __future__ import unicode_literals
+import os
+from csv import reader
+import unittest
+from tempfile import mkdtemp
+from shutil import rmtree
+
+from six import PY3, text_type, binary_type
+
+from dataset import connect
+
+from .sample_data import TEST_DATA
+
+
+class FreezeTestCase(unittest.TestCase):
+
+    def setUp(self):
+        self.db = connect('sqlite://')
+        self.tbl = self.db['weather']
+        for row in TEST_DATA:
+            self.tbl.insert(row)
+        self.d = mkdtemp()
+
+    def tearDown(self):
+        rmtree(self.d, ignore_errors=True)
+
+    def test_freeze(self):
+        from dataset.freeze.app import freeze
+
+        freeze(self.db['weather'].all(), format='csv', filename='wäther.csv'.encode('utf8'), prefix=self.d)
+        self.assert_(os.path.exists(os.path.join(self.d, 'wäther.csv')))
+        freeze(self.db['weather'].all(), format='csv', filename='wäther.csv', prefix=self.d)
+        self.assert_(os.path.exists(os.path.join(self.d, 'wäther.csv')))
+
+    def test_freeze_csv(self):
+        from dataset.freeze.app import freeze
+        from dataset.freeze.format.fcsv import value_to_str
+
+        freeze(self.db['weather'].all(), format='csv', filename='weather.csv', prefix=self.d)
+        path = os.path.join(self.d, 'weather.csv')
+        if PY3:
+            fh = open(path, 'rt', encoding='utf8', newline='')
+        else:
+            fh = open(path, 'rU')
+        rows = list(reader(fh))
+        keys = rows[0]
+        if not PY3:
+            keys = [k.decode('utf8') for k in keys]
+        for i, d1 in enumerate(TEST_DATA):
+            d2 = dict(zip(keys, rows[i + 1]))
+            for k in d1.keys():
+                v2 = d2[k]
+                if not PY3:
+                    v2 = v2.decode('utf8')
+                v1 = value_to_str(d1[k])
+                if not isinstance(v1, text_type):
+                    if isinstance(v1, binary_type):
+                        v1 = text_type(v1, encoding='utf8')
+                    else:
+                        v1 = '%s' % v1
+                self.assertEqual(v2, v1)

test/test_freeze_config.py

@@ -0,0 +1,35 @@
+import os
+import unittest
+
+from dataset.util import FreezeException
+
+
+class TestConfiguration(unittest.TestCase):
+    def test_init(self):
+        from dataset.freeze.config import Configuration
+
+        self.assertRaises(FreezeException, Configuration, 'x.x')
+        self.assertRaises(FreezeException, Configuration, __file__)
+        cfg = Configuration(os.path.join(os.path.dirname(__file__), 'Freezefile.yaml'))
+        assert cfg
+
+    def test_exports(self):
+        from dataset.freeze.config import Configuration
+
+        cfg = Configuration(os.path.join(os.path.dirname(__file__), 'Freezefile.yaml'))
+        exports = list(cfg.exports)
+        self.assertEqual(len(exports), 4)
+        self.assertFalse(exports[0].skip)
+        self.assertTrue(exports[0].get_bool('bool'))
+        self.assertEqual(exports[0].get_int('nan', 'default'), 'default')
+        self.assertEqual(exports[0].get_int('number'), 5)
+        self.assert_(exports[0].name)
+
+    def test_exports_fail(self):
+        from dataset.freeze.config import Configuration
+
+        cfg = Configuration(os.path.join(os.path.dirname(__file__), 'Freezefile.yaml'))
+        cfg.data = None
+        self.assertRaises(FreezeException, list, cfg.exports)
+        cfg.data = {}
+        self.assertRaises(FreezeException, list, cfg.exports)