Add filter `Spaces around Operators`

2 files changed, 39 insertions, 0 deletions

diff --git a/sqlparse/filters.py b/sqlparse/filters.py
index 193029f..464a570 100644
--- a/sqlparse/filters.py
+++ b/sqlparse/filters.py
@@ -139,6 +139,36 @@ class StripWhitespaceFilter(object):
             stmt.tokens.pop(-1)
 
 
+class SpacesAroundOperatorsFilter:
+    whitelist = (sql.Identifier, sql.Comparison, sql.Where)
+
+    def _process(self, tlist):
+        def next_token(idx):
+            # HACK: distinguish between real wildcard from multiplication operator
+            return tlist.token_next_by_type(idx, (T.Operator, T.Comparison, T.Wildcard))
+        idx = 0
+        token = next_token(idx)
+        while token:
+            idx = tlist.token_index(token)
+            if idx > 0 and tlist.tokens[idx - 1].ttype != T.Whitespace:
+                tlist.tokens.insert(idx, sql.Token(T.Whitespace, ' '))  # insert before
+                idx += 1
+            if idx < len(tlist.tokens) - 1:
+                if token.ttype == T.Wildcard and tlist.tokens[idx + 1].match(T.Punctuation, ','):
+                    pass  # this must have been a real wildcard, not multiplication
+                elif tlist.tokens[idx + 1].ttype != T.Whitespace:
+                    tlist.tokens.insert(idx + 1, sql.Token(T.Whitespace, ' '))
+
+            idx += 1
+            token = next_token(idx)
+
+        for sgroup in tlist.get_sublists():
+            self._process(sgroup)
+
+    def process(self, stack, stmt):
+        self._process(stmt)
+
+
 class ReindentFilter(object):
     def __init__(self, width=2, char=' ', line_width=None, wrap_after=0):
         self.width = width
diff --git a/sqlparse/formatter.py b/sqlparse/formatter.py
index 5af8743..0fa563c 100644
--- a/sqlparse/formatter.py
+++ b/sqlparse/formatter.py
@@ -30,6 +30,11 @@ def validate_options(options):
         raise SQLParseError('Invalid value for strip_comments: %r'
                             % strip_comments)
 
+    use_space_around_operators = options.get('use_space_around_operators', False)
+    if use_space_around_operators not in [True, False]:
+        raise SQLParseError('Invalid value for use_space_around_operators: %r'
+                            % use_space_around_operators)
+
     strip_ws = options.get('strip_whitespace', False)
     if strip_ws not in [True, False]:
         raise SQLParseError('Invalid value for strip_whitespace: %r'
@@ -121,6 +126,10 @@ def build_filter_stack(stack, options):
         stack.preprocess.append(filters.TruncateStringFilter(
             width=options['truncate_strings'], char=options['truncate_char']))
 
+    if options.get('use_space_around_operators', False):
+        stack.enable_grouping()
+        stack.stmtprocess.append(filters.SpacesAroundOperatorsFilter())
+
     # After grouping
     if options.get('strip_comments'):
         stack.enable_grouping()