Fixes stdout --sql output in python2

When output_encoding is set, wraps the output buffer into a io.TextIOWrapper class
This means that the output_buffer has to be a io.IOBase instance
in order to work along with the TextIOWrapper

Handles wrapping when the buffer is Python2 stdtout, which has 'file' type
Adds test for this situation

diff --git a/alembic/compat.py b/alembic/compat.py
index 7d0357dccd67a85338f19817ed94f2458e45f9da..fc88933a1b0f505a59a68c6d1300727ce9a7c651 100644 (file)
--- a/alembic/compat.py
+++ b/alembic/compat.py
@@ -1,3 +1,4 @@
+import io
 import sys
 from sqlalchemy import __version__ as sa_version
 
@@ -78,3 +79,10 @@ def with_metaclass(meta, base=object):
     """Create a base class with a metaclass."""
     return meta("%sBase" % meta.__name__, (base,), {})
 ################################################
+
+if py2k:
+    def writable_buffer(file):
+        return io.FileIO(file.fileno(), 'w')
+else:
+    def writable_buffer(file):
+        return file

diff --git a/alembic/migration.py b/alembic/migration.py
index 2b47c919cca3632ce216d2c29ee1435cfe088572..657e60d83e2d6fa26d4ac447cccf0d0abdc85d8a 100644 (file)
--- a/alembic/migration.py
+++ b/alembic/migration.py
@@ -6,7 +6,7 @@ from sqlalchemy import MetaData, Table, Column, String, literal_column
 from sqlalchemy import create_engine
 from sqlalchemy.engine import url as sqla_url
 
-from .compat import callable
+from .compat import callable, writable_buffer
 from . import ddl, util
 
 log = logging.getLogger(__name__)
@@ -71,13 +71,14 @@ class MigrationContext(object):
             self.connection = connection
         self._migrations_fn = opts.get('fn')
         self.as_sql = as_sql
+
         self.output_buffer = opts.get("output_buffer", sys.stdout)
 
-        if opts.get('output_encoding'):
+        if "output_encoding" in opts:
             self.output_buffer = io.TextIOWrapper(
-                                    self.output_buffer,
-                                    opts['output_encoding']
-                                )
+                opts.get("output_buffer") or writable_buffer(sys.stdout),
+                opts.get('output_encoding')
+            )
 
         self._user_compare_type = opts.get('compare_type', False)
         self._user_compare_server_default = opts.get(

diff --git a/tests/test_offline_environment.py b/tests/test_offline_environment.py
index 4a727d369397bdb97f10a41c11a4991d65b7072f..da2589c587cf824a55f720c821587ed275376f60 100644 (file)
--- a/tests/test_offline_environment.py
+++ b/tests/test_offline_environment.py
@@ -1,3 +1,4 @@
+import io
 from unittest import TestCase
 
 from alembic import command, util
@@ -11,6 +12,7 @@ class OfflineEnvironmentTest(TestCase):
     def setUp(self):
         env = staging_env()
         self.cfg = _no_sql_testing_config()
+        self.cfg.output_buffer = io.StringIO()
 
         global a, b, c
         a, b, c = three_rev_fixture(self.cfg)
@@ -41,6 +43,8 @@ assert context.get_starting_revision_argument() == 'x'
         command.upgrade(self.cfg, a, sql=True)
         command.downgrade(self.cfg, "%s:%s" % (b, a), sql=True)
         command.current(self.cfg)
+        # current seems to close the buffer (?)
+        self.cfg.output_buffer = io.StringIO()
         command.stamp(self.cfg, a)
 
     def test_starting_rev_pre_context(self):
@@ -153,3 +157,12 @@ context.configure(dialect_name='sqlite')
             command.downgrade,
             self.cfg, b, sql=True
         )
+
+    def test_upgrade_with_output_encoding(self):
+        env_file_fixture("""
+url = config.get_main_option('sqlalchemy.url')
+context.configure(url=url, output_encoding='utf-8')
+assert not context.requires_connection()
+""")
+        command.upgrade(self.cfg, a, sql=True)
+        command.downgrade(self.cfg, "%s:%s" % (b, a), sql=True)