grouping.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454
  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
  7. from sqlparse import sql
  8. from sqlparse import tokens as T
  9. from sqlparse.utils import recurse, imt
  10. T_NUMERICAL = (T.Number, T.Number.Integer, T.Number.Float)
  11. T_STRING = (T.String, T.String.Single, T.String.Symbol)
  12. T_NAME = (T.Name, T.Name.Placeholder)
  13. def _group_matching(tlist, cls):
  14. """Groups Tokens that have beginning and end."""
  15. opens = []
  16. tidx_offset = 0
  17. for idx, token in enumerate(list(tlist)):
  18. tidx = idx - tidx_offset
  19. if token.is_whitespace:
  20. # ~50% of tokens will be whitespace. Will checking early
  21. # for them avoid 3 comparisons, but then add 1 more comparison
  22. # for the other ~50% of tokens...
  23. continue
  24. if token.is_group and not isinstance(token, cls):
  25. # Check inside previously grouped (i.e. parenthesis) if group
  26. # of different type is inside (i.e., case). though ideally should
  27. # should check for all open/close tokens at once to avoid recursion
  28. _group_matching(token, cls)
  29. continue
  30. if token.match(*cls.M_OPEN):
  31. opens.append(tidx)
  32. elif token.match(*cls.M_CLOSE):
  33. try:
  34. open_idx = opens.pop()
  35. except IndexError:
  36. # this indicates invalid sql and unbalanced tokens.
  37. # instead of break, continue in case other "valid" groups exist
  38. continue
  39. close_idx = tidx
  40. tlist.group_tokens(cls, open_idx, close_idx)
  41. tidx_offset += close_idx - open_idx
  42. def group_brackets(tlist):
  43. _group_matching(tlist, sql.SquareBrackets)
  44. def group_parenthesis(tlist):
  45. _group_matching(tlist, sql.Parenthesis)
  46. def group_case(tlist):
  47. _group_matching(tlist, sql.Case)
  48. def group_if(tlist):
  49. _group_matching(tlist, sql.If)
  50. def group_for(tlist):
  51. _group_matching(tlist, sql.For)
  52. def group_begin(tlist):
  53. _group_matching(tlist, sql.Begin)
  54. def group_typecasts(tlist):
  55. def match(token):
  56. return token.match(T.Punctuation, '::')
  57. def valid(token):
  58. return token is not None
  59. def post(tlist, pidx, tidx, nidx):
  60. return pidx, nidx
  61. valid_prev = valid_next = valid
  62. _group(tlist, sql.Identifier, match, valid_prev, valid_next, post)
  63. def group_tzcasts(tlist):
  64. def match(token):
  65. return token.ttype == T.Keyword.TZCast
  66. def valid(token):
  67. return token is not None
  68. def post(tlist, pidx, tidx, nidx):
  69. return pidx, nidx
  70. _group(tlist, sql.Identifier, match, valid, valid, post)
  71. def group_typed_literal(tlist):
  72. # definitely not complete, see e.g.:
  73. # https://docs.microsoft.com/en-us/sql/odbc/reference/appendixes/interval-literal-syntax
  74. # https://docs.microsoft.com/en-us/sql/odbc/reference/appendixes/interval-literals
  75. # https://www.postgresql.org/docs/9.1/datatype-datetime.html
  76. # https://www.postgresql.org/docs/9.1/functions-datetime.html
  77. def match(token):
  78. return imt(token, m=sql.TypedLiteral.M_OPEN)
  79. def match_to_extend(token):
  80. return isinstance(token, sql.TypedLiteral)
  81. def valid_prev(token):
  82. return token is not None
  83. def valid_next(token):
  84. return token is not None and token.match(*sql.TypedLiteral.M_CLOSE)
  85. def valid_final(token):
  86. return token is not None and token.match(*sql.TypedLiteral.M_EXTEND)
  87. def post(tlist, pidx, tidx, nidx):
  88. return tidx, nidx
  89. _group(tlist, sql.TypedLiteral, match, valid_prev, valid_next,
  90. post, extend=False)
  91. _group(tlist, sql.TypedLiteral, match_to_extend, valid_prev, valid_final,
  92. post, extend=True)
  93. def group_period(tlist):
  94. def match(token):
  95. return token.match(T.Punctuation, '.')
  96. def valid_prev(token):
  97. sqlcls = sql.SquareBrackets, sql.Identifier
  98. ttypes = T.Name, T.String.Symbol
  99. return imt(token, i=sqlcls, t=ttypes)
  100. def valid_next(token):
  101. # issue261, allow invalid next token
  102. return True
  103. def post(tlist, pidx, tidx, nidx):
  104. # next_ validation is being performed here. issue261
  105. sqlcls = sql.SquareBrackets, sql.Function
  106. ttypes = T.Name, T.String.Symbol, T.Wildcard
  107. next_ = tlist[nidx] if nidx is not None else None
  108. valid_next = imt(next_, i=sqlcls, t=ttypes)
  109. return (pidx, nidx) if valid_next else (pidx, tidx)
  110. _group(tlist, sql.Identifier, match, valid_prev, valid_next, post)
  111. def group_as(tlist):
  112. def match(token):
  113. return token.is_keyword and token.normalized == 'AS'
  114. def valid_prev(token):
  115. return token.normalized == 'NULL' or not token.is_keyword
  116. def valid_next(token):
  117. ttypes = T.DML, T.DDL, T.CTE
  118. return not imt(token, t=ttypes) and token is not None
  119. def post(tlist, pidx, tidx, nidx):
  120. return pidx, nidx
  121. _group(tlist, sql.Identifier, match, valid_prev, valid_next, post)
  122. def group_assignment(tlist):
  123. def match(token):
  124. return token.match(T.Assignment, ':=')
  125. def valid(token):
  126. return token is not None and token.ttype not in (T.Keyword)
  127. def post(tlist, pidx, tidx, nidx):
  128. m_semicolon = T.Punctuation, ';'
  129. snidx, _ = tlist.token_next_by(m=m_semicolon, idx=nidx)
  130. nidx = snidx or nidx
  131. return pidx, nidx
  132. valid_prev = valid_next = valid
  133. _group(tlist, sql.Assignment, match, valid_prev, valid_next, post)
  134. def group_comparison(tlist):
  135. sqlcls = (sql.Parenthesis, sql.Function, sql.Identifier,
  136. sql.Operation, sql.TypedLiteral)
  137. ttypes = T_NUMERICAL + T_STRING + T_NAME
  138. def match(token):
  139. return token.ttype == T.Operator.Comparison
  140. def valid(token):
  141. if imt(token, t=ttypes, i=sqlcls):
  142. return True
  143. elif token and token.is_keyword and token.normalized == 'NULL':
  144. return True
  145. else:
  146. return False
  147. def post(tlist, pidx, tidx, nidx):
  148. return pidx, nidx
  149. valid_prev = valid_next = valid
  150. _group(tlist, sql.Comparison, match,
  151. valid_prev, valid_next, post, extend=False)
  152. @recurse(sql.Identifier)
  153. def group_identifier(tlist):
  154. ttypes = (T.String.Symbol, T.Name)
  155. tidx, token = tlist.token_next_by(t=ttypes)
  156. while token:
  157. tlist.group_tokens(sql.Identifier, tidx, tidx)
  158. tidx, token = tlist.token_next_by(t=ttypes, idx=tidx)
  159. def group_arrays(tlist):
  160. sqlcls = sql.SquareBrackets, sql.Identifier, sql.Function
  161. ttypes = T.Name, T.String.Symbol
  162. def match(token):
  163. return isinstance(token, sql.SquareBrackets)
  164. def valid_prev(token):
  165. return imt(token, i=sqlcls, t=ttypes)
  166. def valid_next(token):
  167. return True
  168. def post(tlist, pidx, tidx, nidx):
  169. return pidx, tidx
  170. _group(tlist, sql.Identifier, match,
  171. valid_prev, valid_next, post, extend=True, recurse=False)
  172. def group_operator(tlist):
  173. ttypes = T_NUMERICAL + T_STRING + T_NAME
  174. sqlcls = (sql.SquareBrackets, sql.Parenthesis, sql.Function,
  175. sql.Identifier, sql.Operation, sql.TypedLiteral)
  176. def match(token):
  177. return imt(token, t=(T.Operator, T.Wildcard))
  178. def valid(token):
  179. return imt(token, i=sqlcls, t=ttypes) \
  180. or (token and token.match(
  181. T.Keyword,
  182. ('CURRENT_DATE', 'CURRENT_TIME', 'CURRENT_TIMESTAMP')))
  183. def post(tlist, pidx, tidx, nidx):
  184. tlist[tidx].ttype = T.Operator
  185. return pidx, nidx
  186. valid_prev = valid_next = valid
  187. _group(tlist, sql.Operation, match,
  188. valid_prev, valid_next, post, extend=False)
  189. def group_identifier_list(tlist):
  190. m_role = T.Keyword, ('null', 'role')
  191. sqlcls = (sql.Function, sql.Case, sql.Identifier, sql.Comparison,
  192. sql.IdentifierList, sql.Operation)
  193. ttypes = (T_NUMERICAL + T_STRING + T_NAME
  194. + (T.Keyword, T.Comment, T.Wildcard))
  195. def match(token):
  196. return token.match(T.Punctuation, ',')
  197. def valid(token):
  198. return imt(token, i=sqlcls, m=m_role, t=ttypes)
  199. def post(tlist, pidx, tidx, nidx):
  200. return pidx, nidx
  201. valid_prev = valid_next = valid
  202. _group(tlist, sql.IdentifierList, match,
  203. valid_prev, valid_next, post, extend=True)
  204. @recurse(sql.Comment)
  205. def group_comments(tlist):
  206. tidx, token = tlist.token_next_by(t=T.Comment)
  207. while token:
  208. eidx, end = tlist.token_not_matching(
  209. lambda tk: imt(tk, t=T.Comment) or tk.is_whitespace, idx=tidx)
  210. if end is not None:
  211. eidx, end = tlist.token_prev(eidx, skip_ws=False)
  212. tlist.group_tokens(sql.Comment, tidx, eidx)
  213. tidx, token = tlist.token_next_by(t=T.Comment, idx=tidx)
  214. @recurse(sql.Where)
  215. def group_where(tlist):
  216. tidx, token = tlist.token_next_by(m=sql.Where.M_OPEN)
  217. while token:
  218. eidx, end = tlist.token_next_by(m=sql.Where.M_CLOSE, idx=tidx)
  219. if end is None:
  220. end = tlist._groupable_tokens[-1]
  221. else:
  222. end = tlist.tokens[eidx - 1]
  223. # TODO: convert this to eidx instead of end token.
  224. # i think above values are len(tlist) and eidx-1
  225. eidx = tlist.token_index(end)
  226. tlist.group_tokens(sql.Where, tidx, eidx)
  227. tidx, token = tlist.token_next_by(m=sql.Where.M_OPEN, idx=tidx)
  228. @recurse()
  229. def group_aliased(tlist):
  230. I_ALIAS = (sql.Parenthesis, sql.Function, sql.Case, sql.Identifier,
  231. sql.Operation, sql.Comparison)
  232. tidx, token = tlist.token_next_by(i=I_ALIAS, t=T.Number)
  233. while token:
  234. nidx, next_ = tlist.token_next(tidx)
  235. if isinstance(next_, sql.Identifier):
  236. tlist.group_tokens(sql.Identifier, tidx, nidx, extend=True)
  237. tidx, token = tlist.token_next_by(i=I_ALIAS, t=T.Number, idx=tidx)
  238. @recurse(sql.Function)
  239. def group_functions(tlist):
  240. has_create = False
  241. has_table = False
  242. for tmp_token in tlist.tokens:
  243. if tmp_token.value == 'CREATE':
  244. has_create = True
  245. if tmp_token.value == 'TABLE':
  246. has_table = True
  247. if has_create and has_table:
  248. return
  249. tidx, token = tlist.token_next_by(t=T.Name)
  250. while token:
  251. nidx, next_ = tlist.token_next(tidx)
  252. if isinstance(next_, sql.Parenthesis):
  253. tlist.group_tokens(sql.Function, tidx, nidx)
  254. tidx, token = tlist.token_next_by(t=T.Name, idx=tidx)
  255. def group_order(tlist):
  256. """Group together Identifier and Asc/Desc token"""
  257. tidx, token = tlist.token_next_by(t=T.Keyword.Order)
  258. while token:
  259. pidx, prev_ = tlist.token_prev(tidx)
  260. if imt(prev_, i=sql.Identifier, t=T.Number):
  261. tlist.group_tokens(sql.Identifier, pidx, tidx)
  262. tidx = pidx
  263. tidx, token = tlist.token_next_by(t=T.Keyword.Order, idx=tidx)
  264. @recurse()
  265. def align_comments(tlist):
  266. tidx, token = tlist.token_next_by(i=sql.Comment)
  267. while token:
  268. pidx, prev_ = tlist.token_prev(tidx)
  269. if isinstance(prev_, sql.TokenList):
  270. tlist.group_tokens(sql.TokenList, pidx, tidx, extend=True)
  271. tidx = pidx
  272. tidx, token = tlist.token_next_by(i=sql.Comment, idx=tidx)
  273. def group_values(tlist):
  274. tidx, token = tlist.token_next_by(m=(T.Keyword, 'VALUES'))
  275. start_idx = tidx
  276. end_idx = -1
  277. while token:
  278. if isinstance(token, sql.Parenthesis):
  279. end_idx = tidx
  280. tidx, token = tlist.token_next(tidx)
  281. if end_idx != -1:
  282. tlist.group_tokens(sql.Values, start_idx, end_idx, extend=True)
  283. def group(stmt):
  284. for func in [
  285. group_comments,
  286. # _group_matching
  287. group_brackets,
  288. group_parenthesis,
  289. group_case,
  290. group_if,
  291. group_for,
  292. group_begin,
  293. group_functions,
  294. group_where,
  295. group_period,
  296. group_arrays,
  297. group_identifier,
  298. group_order,
  299. group_typecasts,
  300. group_tzcasts,
  301. group_typed_literal,
  302. group_operator,
  303. group_comparison,
  304. group_as,
  305. group_aliased,
  306. group_assignment,
  307. align_comments,
  308. group_identifier_list,
  309. group_values,
  310. ]:
  311. func(stmt)
  312. return stmt
  313. def _group(tlist, cls, match,
  314. valid_prev=lambda t: True,
  315. valid_next=lambda t: True,
  316. post=None,
  317. extend=True,
  318. recurse=True
  319. ):
  320. """Groups together tokens that are joined by a middle token. i.e. x < y"""
  321. tidx_offset = 0
  322. pidx, prev_ = None, None
  323. for idx, token in enumerate(list(tlist)):
  324. tidx = idx - tidx_offset
  325. if tidx < 0: # tidx shouldn't get negative
  326. continue
  327. if token.is_whitespace:
  328. continue
  329. if recurse and token.is_group and not isinstance(token, cls):
  330. _group(token, cls, match, valid_prev, valid_next, post, extend)
  331. if match(token):
  332. nidx, next_ = tlist.token_next(tidx)
  333. if prev_ and valid_prev(prev_) and valid_next(next_):
  334. from_idx, to_idx = post(tlist, pidx, tidx, nidx)
  335. grp = tlist.group_tokens(cls, from_idx, to_idx, extend=extend)
  336. tidx_offset += to_idx - from_idx
  337. pidx, prev_ = from_idx, grp
  338. continue
  339. pidx, prev_ = tidx, token