Merge pull request andialbrecht#248 from fimmtiu/add-wrap-after-option

andialbrecht · andialbrecht · commit 12d9aee564f1 · 2016-06-03T19:13:02.000+02:00
Add --wrap_after option for wrapping identifier lists.
diff --git a/bin/sqlformat b/bin/sqlformat
@@ -44,6 +44,8 @@ group.add_option('-r', '--reindent', dest='reindent',
                  help='reindent statements')
 group.add_option('--indent_width', dest='indent_width', default=2,
                  help='indentation width (defaults to 2 spaces)')
+group.add_option('--wrap_after', dest='wrap_after', default=0,
+                 help='Column after which lists should be wrapped')
 
 _FORMATTING_GROUP = group
 
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -52,6 +52,10 @@ The :meth:`~sqlparse.format` function accepts the following keyword arguments.
 ``indent_width``
   The width of the indentation, defaults to 2.
 
+``wrap_after``
+  The column limit for wrapping comma-separated lists. If unspecified, it
+  puts every item in the list on its own line.
+
 ``output_format``
   If given the output is additionally formatted to be used as a variable
   in a programming language. Allowed values are "python" and "php".
diff --git a/docs/sqlformat.1 b/docs/sqlformat.1
@@ -49,6 +49,10 @@ Set indent width to
 .IR INDENT_WIDTH .
 Default is 2 spaces.
 .TP
+\fB\-\-wrap_after\fR=\fIWRAP_AFTER\fR
+The column limit for wrapping comma-separated lists. If unspecified, it
+puts every item in the list on its own line.
+.TP
 \fB\-\-strip\-comments
 Remove comments.
 .TP
diff --git a/sqlparse/filters.py b/sqlparse/filters.py
@@ -279,12 +279,13 @@ def process(self, stack, stmt, depth=0):
 
 class ReindentFilter(object):
 
-    def __init__(self, width=2, char=' ', line_width=None):
+    def __init__(self, width=2, char=' ', line_width=None, wrap_after=0):
         self.width = width
         self.char = char
         self.indent = 0
         self.offset = 0
         self.line_width = line_width
+        self.wrap_after = wrap_after
         self._curr_stmt = None
         self._last_stmt = None
 
@@ -413,8 +414,12 @@ def _process_identifierlist(self, tlist):
             else:
                 num_offset = self._get_offset(first) - len(first.value)
             self.offset += num_offset
+            position = self.offset
             for token in identifiers[1:]:
-                tlist.insert_before(token, self.nl())
+                position += len(token.value) + 1   # Add 1 for the "," separator
+                if position > self.wrap_after:
+                    tlist.insert_before(token, self.nl())
+                    position = self.offset
             self.offset -= num_offset
         self._process_default(tlist)
 
diff --git a/sqlparse/formatter.py b/sqlparse/formatter.py
@@ -52,22 +52,33 @@ def validate_options(options):
                             % reindent)
     elif reindent:
         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)
     elif indent_tabs:
         options['indent_char'] = '\t'
     else:
         options['indent_char'] = ' '
+
     indent_width = options.get('indent_width', 2)
     try:
         indent_width = int(indent_width)
     except (TypeError, ValueError):
         raise SQLParseError('indent_width requires an integer')
     if indent_width < 1:
-        raise SQLParseError('indent_width requires an positive integer')
+        raise SQLParseError('indent_width requires a positive integer')
     options['indent_width'] = indent_width
 
+    wrap_after = options.get('wrap_after', 0)
+    try:
+        wrap_after = int(wrap_after)
+    except (TypeError, ValueError):
+        raise SQLParseError('wrap_after requires an integer')
+    if wrap_after < 0:
+        raise SQLParseError('wrap_after requires a positive integer')
+    options['wrap_after'] = wrap_after
+
     right_margin = options.get('right_margin', None)
     if right_margin is not None:
         try:
@@ -115,7 +126,8 @@ def build_filter_stack(stack, options):
         stack.enable_grouping()
         stack.stmtprocess.append(
             filters.ReindentFilter(char=options['indent_char'],
-                                   width=options['indent_width']))
+                                   width=options['indent_width'],
+                                   wrap_after=options['wrap_after']))
 
     if options.get('right_margin', False):
         stack.enable_grouping()
diff --git a/tests/test_format.py b/tests/test_format.py
@@ -117,6 +117,10 @@ def test_option(self):
                           reindent=True, indent_width='foo')
         self.assertRaises(SQLParseError, sqlparse.format, 'foo',
                           reindent=True, indent_width=-12)
+        self.assertRaises(SQLParseError, sqlparse.format, 'foo',
+                          reindent=True, wrap_after='foo')
+        self.assertRaises(SQLParseError, sqlparse.format, 'foo',
+                          reindent=True, wrap_after=-12)
 
     def test_stmts(self):
         f = lambda sql: sqlparse.format(sql, reindent=True)
@@ -204,6 +208,14 @@ def test_identifier_list(self):
                                                'from a,',
                                                '     b']))
 
+    def test_identifier_list_with_wrap_after(self):
+        f = lambda sql: sqlparse.format(sql, reindent=True, wrap_after=14)
+        s = 'select foo, bar, baz from table1, table2 where 1 = 2'
+        self.ndiffAssertEqual(f(s), '\n'.join(['select foo, bar,',
+                                               '       baz',
+                                               'from table1, table2',
+                                               'where 1 = 2']))
+
     def test_identifier_list_with_functions(self):
         f = lambda sql: sqlparse.format(sql, reindent=True)
         s = ("select 'abc' as foo, coalesce(col1, col2)||col3 as bar,"
