From: Hong Minhee <minhee@dahlia.kr>
Date: Fri, 12 Apr 2013 19:40:54 +0000 (+0900)
Subject: Use io instead of codecs and StringIO
X-Git-Tag: rel_0_6_0~22^2~5
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=ff285f75b0132aecdc8b00b1dfabb3eafce8097a;p=thirdparty%2Fsqlalchemy%2Falembic.git

Use io instead of codecs and StringIO
---

diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py
index 1c44534f..af93349c 100644
--- a/alembic/ddl/impl.py
+++ b/alembic/ddl/impl.py
@@ -54,7 +54,10 @@ class DefaultImpl(ImplMeta('_ImplBase', (object,), {})):
         return _impls[dialect.name]
 
     def static_output(self, text):
-        self.output_buffer.write(text + "\n\n")
+        text_ = getattr(builtins, 'unicode', str)(text + '\n\n')
+        self.output_buffer.write(text_)
+        if callable(getattr(self.output_buffer, 'flush', None)):
+            self.output_buffer.flush()
 
     @property
     def bind(self):
diff --git a/alembic/migration.py b/alembic/migration.py
index 82fee9ca..bf8e9380 100644
--- a/alembic/migration.py
+++ b/alembic/migration.py
@@ -1,4 +1,4 @@
-import codecs
+import io
 import logging
 import sys
 
@@ -71,10 +71,12 @@ class MigrationContext(object):
         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'):
-            self.output_buffer = codecs.getwriter(
+        if (opts.get('output_encoding') and
+            not isinstance(self.output_buffer, io.TextIOBase)):
+            self.output_buffer = io.TextIOWrapper(
+                                    self.output_buffer,
                                     opts['output_encoding']
-                                )(self.output_buffer)
+                                )
 
         self._user_compare_type = opts.get('compare_type', False)
         self._user_compare_server_default = opts.get(
diff --git a/tests/__init__.py b/tests/__init__.py
index b0edf75f..37f24d33 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -12,7 +12,6 @@ import io
 import os
 import re
 import shutil
-import StringIO
 import textwrap
 
 from nose import SkipTest
@@ -105,7 +104,12 @@ def assert_compiled(element, assert_string, dialect=None):
     )
 
 def capture_context_buffer(**kw):
-    buf = StringIO.StringIO()
+    if kw.pop('bytes_io', False):
+        raw = io.BytesIO()
+        encoding = kw.get('output_encoding', 'utf-8')
+        buf = io.TextIOWrapper(raw, encoding)
+    else:
+        raw = buf = io.StringIO()
 
     class capture(object):
         def __enter__(self):
@@ -114,7 +118,7 @@ def capture_context_buffer(**kw):
                 'output_buffer':buf
             }
             EnvironmentContext._default_opts.update(kw)
-            return buf
+            return raw
 
         def __exit__(self, *arg, **kwarg):
             #print(buf.getvalue())
diff --git a/tests/test_sql_script.py b/tests/test_sql_script.py
index 6cea44ce..d3511513 100644
--- a/tests/test_sql_script.py
+++ b/tests/test_sql_script.py
@@ -111,6 +111,9 @@ def downgrade():
         clear_staging_env()
 
     def test_encode(self):
-        with capture_context_buffer(output_encoding='utf-8') as buf:
+        with capture_context_buffer(
+                    bytes_io=True,
+                    output_encoding='utf-8'
+                ) as buf:
             command.upgrade(cfg, a, sql=True)
         assert "« S’il vous plaît…".encode("utf-8") in buf.getvalue()