Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1#
2# Copyright (C) 2009-2020 the sqlparse authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of python-sqlparse and is released under
6# the BSD License: https://opensource.org/licenses/BSD-3-Clause
8from sqlparse import sql
9from sqlparse import tokens as T
10from sqlparse.utils import indent, offset
13class ReindentFilter:
14 def __init__(self, width=2, char=' ', wrap_after=0, n='\n',
15 comma_first=False, indent_after_first=False,
16 indent_columns=False, compact=False):
17 self.n = n
18 self.width = width
19 self.char = char
20 self.indent = 1 if indent_after_first else 0
21 self.offset = 0
22 self.wrap_after = wrap_after
23 self.comma_first = comma_first
24 self.indent_columns = indent_columns
25 self.compact = compact
26 self._curr_stmt = None
27 self._last_stmt = None
28 self._last_func = None
30 @property
31 def leading_ws(self):
32 return self.offset + self.indent * self.width
34 def _current_line_len(self, token):
35 """Returns the width of what's already emitted on *token*'s line.
37 The tokens preceding *token* are visited last one first, so the walk
38 stops at the line break that starts the current line. Rebuilding the
39 statement prefix from its start instead made every caller
40 O(statement), and the callers running once per group (tuple lists,
41 identifier lists) quadratic in the number of groups -- a CPU
42 exhaustion vector (GHSA-cfqr-cjx5-5jcm). The walk is inlined and
43 counts characters rather than collecting them: both matter, since a
44 line without any break still has to be measured token by token.
45 """
46 length = 0
47 node = token
48 while node is not self._curr_stmt and node.parent is not None:
49 parent = node.parent
50 # ``Token`` doesn't implement ``__eq__``, so ``index()`` is an
51 # identity lookup and safe against tokens sharing a value.
52 stack = parent.tokens[:parent.tokens.index(node)]
53 while stack:
54 prev_ = stack.pop()
55 if prev_.is_group:
56 stack.extend(prev_.tokens)
57 continue
59 value = prev_.value
60 size = len(value)
61 if not size:
62 continue
63 lines = value.splitlines()
64 if len(lines) == 1 and len(lines[0]) == size:
65 # No break in here. ``splitlines()`` hands back the value
66 # itself in that case, so this costs a scan but no copy --
67 # which is what keeps a long break-free line affordable.
68 length += size
69 continue
71 # ``value`` holds the break that starts the current line. The
72 # sentinel keeps a trailing break from collapsing, so ``lines``
73 # always has one entry more than the number of breaks.
74 lines = (value + '.').splitlines()
75 tail = len(lines[-1]) - 1 + length
76 if tail:
77 return tail
78 # Nothing but a break to our right, and ``splitlines()`` drops
79 # that empty line -- so the line to measure is the one before.
80 if len(lines) > 2:
81 return len(lines[-2])
82 length = len(lines[0])
83 node = parent
84 return length
86 def _get_offset(self, token):
87 # Now take current offset into account and return relative offset.
88 return self._current_line_len(token) - len(self.char * self.leading_ws)
90 def nl(self, offset=0):
91 return sql.Token(
92 T.Whitespace,
93 self.n + self.char * max(0, self.leading_ws + offset))
95 def _next_token(self, tlist, idx=-1):
96 split_words = ('FROM', 'STRAIGHT_JOIN$', 'JOIN$', 'AND', 'OR',
97 'GROUP BY', 'ORDER BY', 'UNION', 'VALUES',
98 'SET', 'BETWEEN', 'EXCEPT', 'HAVING', 'LIMIT')
99 m_split = T.Keyword, split_words, True
100 tidx, token = tlist.token_next_by(m=m_split, idx=idx)
102 if token and token.normalized == 'BETWEEN':
103 tidx, token = self._next_token(tlist, tidx)
105 if token and token.normalized == 'AND':
106 tidx, token = self._next_token(tlist, tidx)
108 return tidx, token
110 def _split_kwds(self, tlist):
111 tidx, token = self._next_token(tlist)
112 while token:
113 pidx, prev_ = tlist.token_prev(tidx, skip_ws=False)
114 uprev = str(prev_)
116 if prev_ and prev_.is_whitespace:
117 del tlist.tokens[pidx]
118 tidx -= 1
120 if not (uprev.endswith('\n') or uprev.endswith('\r')):
121 tlist.insert_before(tidx, self.nl())
122 tidx += 1
124 tidx, token = self._next_token(tlist, tidx)
126 def _split_statements(self, tlist):
127 ttypes = T.Keyword.DML, T.Keyword.DDL
128 tidx, token = tlist.token_next_by(t=ttypes)
129 while token:
130 pidx, prev_ = tlist.token_prev(tidx, skip_ws=False)
131 if prev_ and prev_.is_whitespace:
132 del tlist.tokens[pidx]
133 tidx -= 1
134 # only break if it's not the first token
135 if prev_:
136 tlist.insert_before(tidx, self.nl())
137 tidx += 1
138 tidx, token = tlist.token_next_by(t=ttypes, idx=tidx)
140 def _process(self, tlist):
141 func_name = f'_process_{type(tlist).__name__}'
142 func = getattr(self, func_name.lower(), self._process_default)
143 func(tlist)
145 def _process_where(self, tlist):
146 tidx, token = tlist.token_next_by(m=(T.Keyword, 'WHERE'))
147 if not token:
148 return
149 # issue121, errors in statement fixed??
150 tlist.insert_before(tidx, self.nl())
151 with indent(self):
152 self._process_default(tlist)
154 def _process_parenthesis(self, tlist):
155 ttypes = T.Keyword.DML, T.Keyword.DDL
156 _, is_dml_dll = tlist.token_next_by(t=ttypes)
157 fidx, first = tlist.token_next_by(m=sql.Parenthesis.M_OPEN)
158 if first is None:
159 return
161 with indent(self, 1 if is_dml_dll else 0):
162 tlist.tokens.insert(0, self.nl()) if is_dml_dll else None
163 with offset(self, self._get_offset(first) + 1):
164 self._process_default(tlist, not is_dml_dll)
166 def _process_function(self, tlist):
167 self._last_func = tlist[0]
168 self._process_default(tlist)
170 def _process_identifierlist(self, tlist):
171 identifiers = list(tlist.get_identifiers())
172 if self.indent_columns:
173 first = next(identifiers[0].flatten())
174 num_offset = 1 if self.char == '\t' else self.width
175 else:
176 first = next(identifiers.pop(0).flatten())
177 num_offset = 1 if self.char == '\t' else self._get_offset(first)
179 if not tlist.within(sql.Function) and not tlist.within(sql.Values):
180 with offset(self, num_offset):
181 position = 0
182 for token in identifiers:
183 # Add 1 for the "," separator
184 position += len(token.value) + 1
185 if position > (self.wrap_after - self.offset):
186 adjust = 0
187 if self.comma_first:
188 adjust = -2
189 _, comma = tlist.token_prev(
190 tlist.token_index(token))
191 if comma is None:
192 continue
193 token = comma
194 tlist.insert_before(token, self.nl(offset=adjust))
195 if self.comma_first:
196 _, ws = tlist.token_next(
197 tlist.token_index(token), skip_ws=False)
198 if (ws is not None
199 and ws.ttype is not T.Text.Whitespace):
200 tlist.insert_after(
201 token, sql.Token(T.Whitespace, ' '))
202 position = 0
203 else:
204 # ensure whitespace
205 for token in tlist:
206 _, next_ws = tlist.token_next(
207 tlist.token_index(token), skip_ws=False)
208 if token.value == ',' and not next_ws.is_whitespace:
209 tlist.insert_after(
210 token, sql.Token(T.Whitespace, ' '))
212 end_at = self.offset + sum(len(i.value) + 1 for i in identifiers)
213 adjusted_offset = 0
214 if (self.wrap_after > 0
215 and end_at > (self.wrap_after - self.offset)
216 and self._last_func):
217 adjusted_offset = -len(self._last_func.value) - 1
219 with offset(self, adjusted_offset), indent(self):
220 if adjusted_offset < 0:
221 tlist.insert_before(identifiers[0], self.nl())
222 position = 0
223 for token in identifiers:
224 # Add 1 for the "," separator
225 position += len(token.value) + 1
226 if (self.wrap_after > 0
227 and position > (self.wrap_after - self.offset)):
228 adjust = 0
229 tlist.insert_before(token, self.nl(offset=adjust))
230 position = 0
231 self._process_default(tlist)
233 def _process_case(self, tlist):
234 iterable = iter(tlist.get_cases())
235 cond, _ = next(iterable)
236 first = next(cond[0].flatten())
238 with offset(self, self._get_offset(tlist[0])):
239 with offset(self, self._get_offset(first)):
240 for cond, value in iterable:
241 str_cond = ''.join(str(x) for x in cond or [])
242 str_value = ''.join(str(x) for x in value)
243 end_pos = self.offset + 1 + len(str_cond) + len(str_value)
244 if (not self.compact and end_pos > self.wrap_after):
245 token = value[0] if cond is None else cond[0]
246 tlist.insert_before(token, self.nl())
248 # Line breaks on group level are done. let's add an offset of
249 # len "when ", "then ", "else "
250 with offset(self, len("WHEN ")):
251 self._process_default(tlist)
252 end_idx, end = tlist.token_next_by(m=sql.Case.M_CLOSE)
253 if end_idx is not None and not self.compact:
254 tlist.insert_before(end_idx, self.nl())
256 def _process_values(self, tlist):
257 tlist.insert_before(0, self.nl())
258 tidx, token = tlist.token_next_by(i=sql.Parenthesis)
259 first_token = token
260 while token:
261 ptidx, ptoken = tlist.token_next_by(m=(T.Punctuation, ','),
262 idx=tidx)
263 if ptoken:
264 if self.comma_first:
265 adjust = -2
266 offset = self._get_offset(first_token) + adjust
267 tlist.insert_before(ptoken, self.nl(offset))
268 else:
269 tlist.insert_after(ptoken,
270 self.nl(self._get_offset(token)))
271 tidx, token = tlist.token_next_by(i=sql.Parenthesis, idx=tidx)
273 def _process_default(self, tlist, stmts=True):
274 self._split_statements(tlist) if stmts else None
275 self._split_kwds(tlist)
276 for sgroup in tlist.get_sublists():
277 self._process(sgroup)
279 def process(self, stmt):
280 self._curr_stmt = stmt
281 self._process(stmt)
283 if self._last_stmt is not None:
284 nl = '\n' if str(self._last_stmt).endswith('\n') else '\n\n'
285 stmt.tokens.insert(0, sql.Token(T.Whitespace, nl))
287 self._last_stmt = stmt
288 return stmt