Add option to truncate long string literals.

andialbrecht · andialbrecht · commit 204b77954eac · 2013-07-09T10:10:33.000+02:00
diff --git a/CHANGES b/CHANGES
@@ -1,3 +1,10 @@
+Development Version
+-------------------
+
+Enhancements
+* New option to truncate long string literals when formatting.
+
+
 Release 0.1.8 (Jun 29, 2013)
 ----------------------------
 
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -35,6 +35,14 @@ The :meth:`~sqlparse.format` function accepts the following keyword arguments.
 ``strip_comments``
   If ``True`` comments are removed from the statements.
 
+``truncate_strings``
+  If ``truncate_strings`` is a positive integer, string literals longer than
+  the given value will be truncated.
+
+``truncate_char`` (default: "[...]")
+  If long string literals are truncated (see above) this value will be append
+  to the truncated string.
+
 ``reindent``
   If ``True`` the indentations of the statements are changed.
 
diff --git a/sqlparse/filters.py b/sqlparse/filters.py
@@ -47,6 +47,27 @@ def process(self, stack, stream):
             yield ttype, value
 
 
+class TruncateStringFilter:
+
+    def __init__(self, width, char):
+        self.width = max(width, 1)
+        self.char = unicode(char)
+
+    def process(self, stack, stream):
+        for ttype, value in stream:
+            if ttype is T.Literal.String.Single:
+                if value[:2] == '\'\'':
+                    inner = value[2:-2]
+                    quote = u'\'\''
+                else:
+                    inner = value[1:-1]
+                    quote = u'\''
+                if len(inner) > self.width:
+                    value = u''.join((quote, inner[:self.width], self.char,
+                                      quote))
+            yield ttype, value
+
+
 class GetComments:
     """Get the comments from a stack"""
     def process(self, stack, stream):
diff --git a/sqlparse/formatter.py b/sqlparse/formatter.py
@@ -33,6 +33,19 @@ def validate_options(options):
         raise SQLParseError('Invalid value for strip_whitespace: %r'
                             % strip_ws)
 
+    truncate_strings = options.get('truncate_strings', None)
+    if truncate_strings is not None:
+        try:
+            truncate_strings = int(truncate_strings)
+        except (ValueError, TypeError):
+            raise SQLParseError('Invalid value for truncate_strings: %r'
+                                % truncate_strings)
+        if truncate_strings <= 1:
+            raise SQLParseError('Invalid value for truncate_strings: %r'
+                                % truncate_strings)
+        options['truncate_strings'] = truncate_strings
+        options['truncate_char'] = options.get('truncate_char', '[...]')
+
     reindent = options.get('reindent', False)
     if reindent not in [True, False]:
         raise SQLParseError('Invalid value for reindent: %r'
@@ -84,6 +97,10 @@ def build_filter_stack(stack, options):
         stack.preprocess.append(
             filters.IdentifierCaseFilter(options['identifier_case']))
 
+    if options.get('truncate_strings', None) is not None:
+        stack.preprocess.append(filters.TruncateStringFilter(
+            width=options['truncate_strings'], char=options['truncate_char']))
+
     # After grouping
     if options.get('strip_comments', False):
         stack.enable_grouping()
diff --git a/tests/test_format.py b/tests/test_format.py
@@ -1,5 +1,7 @@
 # -*- coding: utf-8 -*-
 
+import pytest
+
 from tests.utils import TestCaseBase
 
 import sqlparse
@@ -283,3 +285,27 @@ def test_format_column_ordering():  # issue89
                           '         c2,',
                           '         c3;'])
     assert formatted == expected
+
+
+def test_truncate_strings():
+    sql = 'update foo set value = \'' + 'x' * 1000 + '\';'
+    formatted = sqlparse.format(sql, truncate_strings=10)
+    assert formatted == 'update foo set value = \'xxxxxxxxxx[...]\';'
+    formatted = sqlparse.format(sql, truncate_strings=3, truncate_char='YYY')
+    assert formatted == 'update foo set value = \'xxxYYY\';'
+
+
+def test_truncate_strings_invalid_option():
+    pytest.raises(SQLParseError, sqlparse.format,
+                  'foo', truncate_strings='bar')
+    pytest.raises(SQLParseError, sqlparse.format,
+                  'foo', truncate_strings=-1)
+    pytest.raises(SQLParseError, sqlparse.format,
+                  'foo', truncate_strings=0)
+
+
+@pytest.mark.parametrize('sql', ['select verrrylongcolumn from foo',
+                                 'select "verrrylongcolumn" from "foo"'])
+def test_truncate_strings_doesnt_truncate_identifiers(sql):
+    formatted = sqlparse.format(sql, truncate_strings=2)
+    assert formatted == sql
