Start work on Aligned-reindent

3 files changed, 136 insertions, 0 deletions

diff --git a/sqlparse/filters.py b/sqlparse/filters.py
index 95ac74c..78c8d34 100644
--- a/sqlparse/filters.py
+++ b/sqlparse/filters.py
@@ -335,6 +335,98 @@ class ReindentFilter(object):
                 self._last_stmt = stmt
 
 
+class AlignedIndentFilter:
+    split_words = (
+        'FROM',
+        'JOIN', 'ON',
+        'WHERE', 'AND', 'OR',
+        'GROUP', 'HAVING',
+        'ORDER', 'UNION', 'VALUES',
+        'SET', 'BETWEEN', 'EXCEPT',
+        )
+
+    def __init__(self, char=' ', line_width=None):
+        self.char = char
+        self._max_kwd_len = len('select')
+
+    def newline(self):
+        return sql.Token(T.Newline, '\n')
+
+    def whitespace(self, chars=0, newline_before=False, newline_after=False):
+        return sql.Token(
+            T.Whitespace,
+            (str(self.newline()) if newline_before else '') + self.char * chars + (str(self.newline()) if newline_after else ''))
+
+    def _process_statement(self, tlist, base_indent=0):
+        if tlist.tokens[0].is_whitespace() and base_indent == 0:
+            tlist.tokens.pop(0)
+
+        # process the main query body
+        return self._process(sql.TokenList(tlist.tokens), base_indent=base_indent)
+
+    def _process_parenthesis(self, tlist, base_indent=0):
+        sub_indent = base_indent + self._max_kwd_len + 2  # add two for the space and parens
+        tlist.insert_after(tlist.tokens[0], self.whitespace(sub_indent, newline_before=True))
+        # de-indent the last parenthesis
+        tlist.insert_before(tlist.tokens[-1], self.whitespace(sub_indent - 1, newline_before=True))
+
+        # process the inside of the parantheses
+        tlist.tokens = (
+            [tlist.tokens[0]] +
+            self._process(sql.TokenList(tlist._groupable_tokens), base_indent=sub_indent).tokens +
+            [tlist.tokens[-1]]
+            )
+        return tlist
+
+    def _process_identifierlist(self, tlist, base_indent=0):
+        # columns being selected
+        new_tokens = []
+        identifiers = filter(lambda t: isinstance(t, sql.Identifier), tlist.tokens)
+        for i, token in enumerate(identifiers):
+            if i > 0:
+                new_tokens.append(self.newline())
+                new_tokens.append(self.whitespace(self._max_kwd_len + base_indent + 1))
+            new_tokens.append(token)
+            if i < len(identifiers) - 1:
+                # if not last column in select, add a comma seperator
+                new_tokens.append(sql.Token(T.Punctuation, ','))
+        tlist.tokens = new_tokens
+        return tlist
+
+    def _process_substatement(self, tlist, base_indent=0):
+        def _next_token(i):
+            t = tlist.token_next_match(i, T.Keyword, self.split_words, regex=True)
+            # treat "BETWEEN x and y" as a single statement
+            if t and t.value.upper() == 'BETWEEN':
+                t = _next_token(tlist.token_index(t) + 1)
+                if t and t.value.upper() == 'AND':
+                    t = _next_token(tlist.token_index(t) + 1)
+            return t
+
+        idx = 0
+        token = _next_token(idx)
+        while token:
+            tlist.insert_before(token, self.whitespace(self._max_kwd_len - len(str(token)) + base_indent, newline_before=True))
+            next_idx = tlist.token_index(token) + 1
+            token = _next_token(next_idx)
+
+        # process any sub-sub statements
+        for sgroup in tlist.get_sublists():
+            self._process(sgroup, base_indent=base_indent)
+        return tlist
+
+    def _process(self, tlist, base_indent=0, verbose=False):
+        token_name = tlist.__class__.__name__.lower()
+        func_name = '_process_%s' % token_name
+        func = getattr(self, func_name, self._process_substatement)
+        if verbose:
+            print func.__name__, token_name, str(tlist)
+        return func(tlist, base_indent=base_indent)
+
+    def process(self, stack, stmt):
+        self._process(stmt)
+
+
 # FIXME: Doesn't work
 class RightMarginFilter(object):
     keep_together = (
diff --git a/sqlparse/formatter.py b/sqlparse/formatter.py
index 7441313..5af8743 100644
--- a/sqlparse/formatter.py
+++ b/sqlparse/formatter.py
@@ -55,6 +55,13 @@ def validate_options(options):
     elif reindent:
         options['strip_whitespace'] = True
 
+    reindent_aligned = options.get('reindent_aligned', False)
+    if reindent_aligned not in [True, False]:
+        raise SQLParseError('Invalid value for reindent_aligned: %r'
+                            % reindent)
+    elif reindent_aligned:
+        options['strip_whitespace'] = True
+
     indent_tabs = options.get('indent_tabs', False)
     if indent_tabs not in [True, False]:
         raise SQLParseError('Invalid value for indent_tabs: %r' % indent_tabs)
@@ -130,6 +137,11 @@ def build_filter_stack(stack, options):
                                    width=options['indent_width'],
                                    wrap_after=options['wrap_after']))
 
+    if options.get('reindent_aligned', False):
+        stack.enable_grouping()
+        stack.stmtprocess.append(
+            filters.AlignedIndentFilter(char=options['indent_char']))
+
     if options.get('right_margin'):
         stack.enable_grouping()
         stack.stmtprocess.append(
diff --git a/tests/test_format.py b/tests/test_format.py
index 9043e76..2f9ee9f 100644
--- a/tests/test_format.py
+++ b/tests/test_format.py
@@ -106,6 +106,38 @@ class TestFormat(TestCaseBase):
                           output_format='foo')
 
 
+class TestFormatReindentAligned(TestCaseBase):
+    @staticmethod
+    def formatter(sql):
+        return sqlparse.format(sql, reindent_aligned=True)
+
+    def test_basic(self):
+        sql = """
+            select a, b as bb,c from table
+            join (select a * 2 as a from new_table) other
+            on table.a = other.a
+            where c is true
+            and b between 3 and 4
+            or d is 'blue'
+            """
+        self.ndiffAssertEqual(
+            self.formatter(sql),
+            '\n'.join([
+                'select a,',
+                '       b as bb,',
+                '       c',
+                '  from table',
+                '  join (',
+                '        select a * 2 as a',
+                '          from new_table',
+                '       ) other',
+                '    on table.a = other.a',
+                ' where c is true',
+                '   and b between 3 and 4',
+                "    or d is 'blue'",
+            ]))
+
+
 class TestFormatReindent(TestCaseBase):
 
     def test_option(self):