utils.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121
  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. import itertools
  8. import re
  9. from collections import deque
  10. from contextlib import contextmanager
  11. # This regular expression replaces the home-cooked parser that was here before.
  12. # It is much faster, but requires an extra post-processing step to get the
  13. # desired results (that are compatible with what you would expect from the
  14. # str.splitlines() method).
  15. #
  16. # It matches groups of characters: newlines, quoted strings, or unquoted text,
  17. # and splits on that basis. The post-processing step puts those back together
  18. # into the actual lines of SQL.
  19. SPLIT_REGEX = re.compile(r"""
  20. (
  21. (?: # Start of non-capturing group
  22. (?:\r\n|\r|\n) | # Match any single newline, or
  23. [^\r\n'"]+ | # Match any character series without quotes or
  24. # newlines, or
  25. "(?:[^"\\]|\\.)*" | # Match double-quoted strings, or
  26. '(?:[^'\\]|\\.)*' # Match single quoted strings
  27. )
  28. )
  29. """, re.VERBOSE)
  30. LINE_MATCH = re.compile(r'(\r\n|\r|\n)')
  31. def split_unquoted_newlines(stmt):
  32. """Split a string on all unquoted newlines.
  33. Unlike str.splitlines(), this will ignore CR/LF/CR+LF if the requisite
  34. character is inside of a string."""
  35. text = str(stmt)
  36. lines = SPLIT_REGEX.split(text)
  37. outputlines = ['']
  38. for line in lines:
  39. if not line:
  40. continue
  41. elif LINE_MATCH.match(line):
  42. outputlines.append('')
  43. else:
  44. outputlines[-1] += line
  45. return outputlines
  46. def remove_quotes(val):
  47. """Helper that removes surrounding quotes from strings."""
  48. if val is None:
  49. return
  50. if val[0] in ('"', "'") and val[0] == val[-1]:
  51. val = val[1:-1]
  52. return val
  53. def recurse(*cls):
  54. """Function decorator to help with recursion
  55. :param cls: Classes to not recurse over
  56. :return: function
  57. """
  58. def wrap(f):
  59. def wrapped_f(tlist):
  60. for sgroup in tlist.get_sublists():
  61. if not isinstance(sgroup, cls):
  62. wrapped_f(sgroup)
  63. f(tlist)
  64. return wrapped_f
  65. return wrap
  66. def imt(token, i=None, m=None, t=None):
  67. """Helper function to simplify comparisons Instance, Match and TokenType
  68. :param token:
  69. :param i: Class or Tuple/List of Classes
  70. :param m: Tuple of TokenType & Value. Can be list of Tuple for multiple
  71. :param t: TokenType or Tuple/List of TokenTypes
  72. :return: bool
  73. """
  74. clss = i
  75. types = [t, ] if t and not isinstance(t, list) else t
  76. mpatterns = [m, ] if m and not isinstance(m, list) else m
  77. if token is None:
  78. return False
  79. elif clss and isinstance(token, clss):
  80. return True
  81. elif mpatterns and any(token.match(*pattern) for pattern in mpatterns):
  82. return True
  83. elif types and any(token.ttype in ttype for ttype in types):
  84. return True
  85. else:
  86. return False
  87. def consume(iterator, n):
  88. """Advance the iterator n-steps ahead. If n is none, consume entirely."""
  89. deque(itertools.islice(iterator, n), maxlen=0)
  90. @contextmanager
  91. def offset(filter_, n=0):
  92. filter_.offset += n
  93. yield
  94. filter_.offset -= n
  95. @contextmanager
  96. def indent(filter_, n=1):
  97. filter_.indent += n
  98. yield
  99. filter_.indent -= n