Development of an internal social media platform with personalised dashboards for students
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

aligned_indent.py 4.8KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129
  1. # -*- coding: utf-8 -*-
  2. #
  3. # Copyright (C) 2016 Andi Albrecht, albrecht.andi@gmail.com
  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, tokens as T
  8. from sqlparse.compat import text_type
  9. from sqlparse.utils import offset, indent
  10. class AlignedIndentFilter(object):
  11. join_words = (r'((LEFT\s+|RIGHT\s+|FULL\s+)?'
  12. r'(INNER\s+|OUTER\s+|STRAIGHT\s+)?|'
  13. r'(CROSS\s+|NATURAL\s+)?)?JOIN\b')
  14. split_words = ('FROM',
  15. join_words, 'ON',
  16. 'WHERE', 'AND', 'OR',
  17. 'GROUP', 'HAVING', 'LIMIT',
  18. 'ORDER', 'UNION', 'VALUES',
  19. 'SET', 'BETWEEN', 'EXCEPT')
  20. def __init__(self, char=' ', n='\n'):
  21. self.n = n
  22. self.offset = 0
  23. self.indent = 0
  24. self.char = char
  25. self._max_kwd_len = len('select')
  26. def nl(self, offset=1):
  27. # offset = 1 represent a single space after SELECT
  28. offset = -len(offset) if not isinstance(offset, int) else offset
  29. # add two for the space and parens
  30. indent = self.indent * (2 + self._max_kwd_len)
  31. return sql.Token(T.Whitespace, self.n + self.char * (
  32. self._max_kwd_len + offset + indent + self.offset))
  33. def _process_statement(self, tlist):
  34. if tlist.tokens[0].is_whitespace and self.indent == 0:
  35. tlist.tokens.pop(0)
  36. # process the main query body
  37. self._process(sql.TokenList(tlist.tokens))
  38. def _process_parenthesis(self, tlist):
  39. # if this isn't a subquery, don't re-indent
  40. _, token = tlist.token_next_by(m=(T.DML, 'SELECT'))
  41. if token is not None:
  42. with indent(self):
  43. tlist.insert_after(tlist[0], self.nl('SELECT'))
  44. # process the inside of the parantheses
  45. self._process_default(tlist)
  46. # de-indent last parenthesis
  47. tlist.insert_before(tlist[-1], self.nl())
  48. def _process_identifierlist(self, tlist):
  49. # columns being selected
  50. identifiers = list(tlist.get_identifiers())
  51. identifiers.pop(0)
  52. [tlist.insert_before(token, self.nl()) for token in identifiers]
  53. self._process_default(tlist)
  54. def _process_case(self, tlist):
  55. offset_ = len('case ') + len('when ')
  56. cases = tlist.get_cases(skip_ws=True)
  57. # align the end as well
  58. end_token = tlist.token_next_by(m=(T.Keyword, 'END'))[1]
  59. cases.append((None, [end_token]))
  60. condition_width = [len(' '.join(map(text_type, cond))) if cond else 0
  61. for cond, _ in cases]
  62. max_cond_width = max(condition_width)
  63. for i, (cond, value) in enumerate(cases):
  64. # cond is None when 'else or end'
  65. stmt = cond[0] if cond else value[0]
  66. if i > 0:
  67. tlist.insert_before(stmt, self.nl(
  68. offset_ - len(text_type(stmt))))
  69. if cond:
  70. ws = sql.Token(T.Whitespace, self.char * (
  71. max_cond_width - condition_width[i]))
  72. tlist.insert_after(cond[-1], ws)
  73. def _next_token(self, tlist, idx=-1):
  74. split_words = T.Keyword, self.split_words, True
  75. tidx, token = tlist.token_next_by(m=split_words, idx=idx)
  76. # treat "BETWEEN x and y" as a single statement
  77. if token and token.normalized == 'BETWEEN':
  78. tidx, token = self._next_token(tlist, tidx)
  79. if token and token.normalized == 'AND':
  80. tidx, token = self._next_token(tlist, tidx)
  81. return tidx, token
  82. def _split_kwds(self, tlist):
  83. tidx, token = self._next_token(tlist)
  84. while token:
  85. # joins are special case. only consider the first word as aligner
  86. if token.match(T.Keyword, self.join_words, regex=True):
  87. token_indent = token.value.split()[0]
  88. else:
  89. token_indent = text_type(token)
  90. tlist.insert_before(token, self.nl(token_indent))
  91. tidx += 1
  92. tidx, token = self._next_token(tlist, tidx)
  93. def _process_default(self, tlist):
  94. self._split_kwds(tlist)
  95. # process any sub-sub statements
  96. for sgroup in tlist.get_sublists():
  97. idx = tlist.token_index(sgroup)
  98. pidx, prev_ = tlist.token_prev(idx)
  99. # HACK: make "group/order by" work. Longer than max_len.
  100. offset_ = 3 if (prev_ and prev_.match(T.Keyword, 'BY')) else 0
  101. with offset(self, offset_):
  102. self._process(sgroup)
  103. def _process(self, tlist):
  104. func_name = '_process_{cls}'.format(cls=type(tlist).__name__)
  105. func = getattr(self, func_name.lower(), self._process_default)
  106. func(tlist)
  107. def process(self, stmt):
  108. self._process(stmt)
  109. return stmt