From 94024c2d2d66d5fbe875d90ec722cac4cf1601f7 Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Thu, 6 Nov 2014 18:49:24 -0500
Subject: [PATCH] - use pyprinter for autogen so that we get indentation logic

---
 alembic/autogenerate/api.py | 64 ++++++++++++++++++++++++-------------
 alembic/compat.py           |  6 ++++
 2 files changed, 47 insertions(+), 23 deletions(-)

diff --git a/alembic/autogenerate/api.py b/alembic/autogenerate/api.py
index 65ea3ec5..3519a6d1 100644
--- a/alembic/autogenerate/api.py
+++ b/alembic/autogenerate/api.py
@@ -4,6 +4,9 @@ automatically."""
 import logging
 import re
 
+from ..compat import StringIO
+
+from mako.pygen import PythonPrinter
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.util import OrderedSet
 from .compare import _compare_tables
@@ -138,13 +141,37 @@ def _produce_migration_diffs(context, template_args,
     diffs = []
     _produce_net_changes(connection, metadata, diffs,
                          autogen_context, object_filters, include_schemas)
-    template_args[opts['upgrade_token']] = \
-        _indent(_produce_upgrade_commands(diffs, autogen_context))
-    template_args[opts['downgrade_token']] = \
-        _indent(_produce_downgrade_commands(diffs, autogen_context))
+
+    template_args[opts['upgrade_token']] = _indent(_render_cmd_body(
+        _produce_upgrade_commands, diffs, autogen_context))
+    template_args[opts['downgrade_token']] = _indent(_render_cmd_body(
+        _produce_downgrade_commands, diffs, autogen_context))
     template_args['imports'] = "\n".join(sorted(imports))
 
 
+def _indent(text):
+    text = re.compile(r'^', re.M).sub("    ", text).strip()
+    return text
+
+
+def _render_cmd_body(fn, diffs, autogen_context):
+
+    buf = StringIO()
+    printer = PythonPrinter(buf)
+
+    printer.writeline(
+        "### commands auto generated by Alembic - "
+        "please adjust! ###"
+    )
+
+    for line in fn(diffs, autogen_context):
+        printer.writeline(line)
+
+    printer.writeline("### end Alembic commands ###")
+
+    return buf.getvalue()
+
+
 def _get_object_filters(
         context_opts, include_symbol=None, include_object=None):
     include_symbol = context_opts.get('include_symbol', include_symbol)
@@ -176,13 +203,6 @@ def _autogen_context(context, imports):
     }, connection
 
 
-def _indent(text):
-    text = "### commands auto generated by Alembic - "\
-        "please adjust! ###\n" + text
-    text += "\n### end Alembic commands ###"
-    text = re.compile(r'^', re.M).sub("    ", text).strip()
-    return text
-
 ###################################################
 # walk structures
 
@@ -225,21 +245,19 @@ def _produce_net_changes(connection, metadata, diffs, autogen_context,
 
 
 def _produce_upgrade_commands(diffs, autogen_context):
-    buf = []
-    for diff in diffs:
-        buf.append(_invoke_command("upgrade", diff, autogen_context))
-    if not buf:
-        buf = ["pass"]
-    return "\n".join(buf)
+    if diffs:
+        for diff in diffs:
+            yield _invoke_command("upgrade", diff, autogen_context)
+    else:
+        yield "pass"
 
 
 def _produce_downgrade_commands(diffs, autogen_context):
-    buf = []
-    for diff in reversed(diffs):
-        buf.append(_invoke_command("downgrade", diff, autogen_context))
-    if not buf:
-        buf = ["pass"]
-    return "\n".join(buf)
+    if diffs:
+        for diff in reversed(diffs):
+            yield _invoke_command("downgrade", diff, autogen_context)
+    else:
+        yield "pass"
 
 
 def _invoke_command(updown, args, autogen_context):
diff --git a/alembic/compat.py b/alembic/compat.py
index c47ce814..8c6a038c 100644
--- a/alembic/compat.py
+++ b/alembic/compat.py
@@ -12,6 +12,12 @@ py2k = sys.version_info < (3, 0)
 py3k = sys.version_info >= (3, 0)
 py33 = sys.version_info >= (3, 3)
 
+if py3k:
+    from io import StringIO
+else:
+    # accepts strings
+    from StringIO import StringIO
+
 if py3k:
     import builtins as compat_builtins
     string_types = str,
-- 
2.47.2