-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathfilter_stack.py
More file actions
45 lines (36 loc) · 1.27 KB
/
filter_stack.py
File metadata and controls
45 lines (36 loc) · 1.27 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
"""filter"""
from sqlparse import lexer
from sqlparse.engine import grouping
from sqlparse.engine.statement_splitter import StatementSplitter
from sqlparse.filters import StripTrailingSemicolonFilter
class FilterStack:
def __init__(self, strip_semicolon=False):
self.preprocess = []
self.stmtprocess = []
self.postprocess = []
self._grouping = False
if strip_semicolon:
self.stmtprocess.append(StripTrailingSemicolonFilter())
def enable_grouping(self):
self._grouping = True
def run(self, sql, encoding=None):
stream = lexer.tokenize(sql, encoding)
# Process token stream
for filter_ in self.preprocess:
stream = filter_.process(stream)
stream = StatementSplitter().process(stream)
# Group and ungroup tokens
if self._grouping:
stream = grouping.group(stream)
# Process statements
ret = []
for stmt in stream:
if stmt.is_whitespace:
continue
for filter_ in self.stmtprocess:
filter_.process(stmt)
ret.append(stmt)
# Process again after grouping
for filter_ in self.postprocess:
ret = filter_.process(ret)
return ret