Merge pull request #398 from intellihr/wrap_long_function

Attempt to wrap long function parameters with wrap_after

2 files changed, 46 insertions, 0 deletions

diff --git a/sqlparse/filters/reindent.py b/sqlparse/filters/reindent.py
index 49c7807..3088d6f 100644
--- a/sqlparse/filters/reindent.py
+++ b/sqlparse/filters/reindent.py
@@ -24,6 +24,7 @@ class ReindentFilter(object):
         self.indent_columns = indent_columns
         self._curr_stmt = None
         self._last_stmt = None
+        self._last_func = None
 
     def _flatten_up_to_token(self, token):
         """Yields all tokens up to token but excluding current."""
@@ -118,6 +119,10 @@ class ReindentFilter(object):
             with offset(self, self._get_offset(first) + 1):
                 self._process_default(tlist, not is_dml_dll)
 
+    def _process_function(self, tlist):
+        self._last_func = tlist[0]
+        self._process_default(tlist)
+
     def _process_identifierlist(self, tlist):
         identifiers = list(tlist.get_identifiers())
         if self.indent_columns:
@@ -126,6 +131,7 @@ class ReindentFilter(object):
         else:
             first = next(identifiers.pop(0).flatten())
             num_offset = 1 if self.char == '\t' else self._get_offset(first)
+
         if not tlist.within(sql.Function):
             with offset(self, num_offset):
                 position = 0
@@ -150,6 +156,34 @@ class ReindentFilter(object):
                                 tlist.insert_after(
                                     token, sql.Token(T.Whitespace, ' '))
                         position = 0
+        else:
+            # ensure whitespace
+            for token in tlist:
+                _, next_ws = tlist.token_next(
+                    tlist.token_index(token), skip_ws=False)
+                if token.value == ',' and not next_ws.is_whitespace:
+                    tlist.insert_after(
+                        token, sql.Token(T.Whitespace, ' '))
+
+            end_at = self.offset + sum(len(i.value) + 1 for i in identifiers)
+            adjusted_offset = 0
+            if (self.wrap_after > 0
+                    and end_at > (self.wrap_after - self.offset)
+                    and self._last_func):
+                adjusted_offset = -len(self._last_func.value) - 1
+
+            with offset(self, adjusted_offset), indent(self):
+                if adjusted_offset < 0:
+                    tlist.insert_before(identifiers[0], self.nl())
+                position = 0
+                for token in identifiers:
+                    # Add 1 for the "," separator
+                    position += len(token.value) + 1
+                    if (self.wrap_after > 0
+                            and position > (self.wrap_after - self.offset)):
+                        adjust = 0
+                        tlist.insert_before(token, self.nl(offset=adjust))
+                        position = 0
         self._process_default(tlist)
 
     def _process_case(self, tlist):
diff --git a/tests/test_format.py b/tests/test_format.py
index 342fe09..72af62e 100644
--- a/tests/test_format.py
+++ b/tests/test_format.py
@@ -453,6 +453,18 @@ class TestFormatReindent(object):
             "       col3",
             "from my_table"])
 
+    def test_long_identifier_list_with_functions(self):
+        f = lambda sql: sqlparse.format(sql, reindent=True, wrap_after=30)
+        s = ("select 'abc' as foo, json_build_object('a',  a,"
+             "'b', b, 'c', c, 'd', d, 'e', e) as col2"
+             "col3 from my_table")
+        assert f(s) == '\n'.join([
+            "select 'abc' as foo,",
+            "       json_build_object('a',",
+            "         a, 'b', b, 'c', c, 'd', d,",
+            "         'e', e) as col2col3",
+            "from my_table"])
+
     def test_case(self):
         f = lambda sql: sqlparse.format(sql, reindent=True)
         s = 'case when foo = 1 then 2 when foo = 3 then 4 else 5 end'