cli.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246
  1. # Copyright (C) 2009-2020 the sqlparse authors and contributors
  2. # <see AUTHORS file>
  3. #
  4. # This module is part of python-sqlparse and is released under
  5. # the BSD License: https://opensource.org/licenses/BSD-3-Clause
  6. """Module that contains the command line app.
  7. Why does this file exist, and why not put this in __main__?
  8. You might be tempted to import things from __main__ later, but that will
  9. cause problems: the code will get executed twice:
  10. - When you run `python -m sqlparse` python will execute
  11. ``__main__.py`` as a script. That means there won't be any
  12. ``sqlparse.__main__`` in ``sys.modules``.
  13. - When you import __main__ it will get executed again (as a module) because
  14. there's no ``sqlparse.__main__`` in ``sys.modules``.
  15. Also see (1) from http://click.pocoo.org/5/setuptools/#setuptools-integration
  16. """
  17. import argparse
  18. import sys
  19. from io import TextIOWrapper
  20. import sqlparse
  21. from sqlparse.exceptions import SQLParseError
  22. # TODO: Add CLI Tests
  23. # TODO: Simplify formatter by using argparse `type` arguments
  24. def create_parser():
  25. _CASE_CHOICES = ['upper', 'lower', 'capitalize']
  26. parser = argparse.ArgumentParser(
  27. prog='sqlformat',
  28. description='Format FILE according to OPTIONS. Use "-" as FILE '
  29. 'to read from stdin.',
  30. usage='%(prog)s [OPTIONS] FILE [FILE ...]',
  31. )
  32. parser.add_argument(
  33. 'filename',
  34. nargs='+',
  35. help='file(s) to format (use "-" for stdin)')
  36. parser.add_argument(
  37. '-o', '--outfile',
  38. dest='outfile',
  39. metavar='FILE',
  40. help='write output to FILE (defaults to stdout)')
  41. parser.add_argument(
  42. '--in-place',
  43. dest='inplace',
  44. action='store_true',
  45. default=False,
  46. help='format files in-place (overwrite existing files)')
  47. parser.add_argument(
  48. '--version',
  49. action='version',
  50. version=sqlparse.__version__)
  51. group = parser.add_argument_group('Formatting Options')
  52. group.add_argument(
  53. '-k', '--keywords',
  54. metavar='CHOICE',
  55. dest='keyword_case',
  56. choices=_CASE_CHOICES,
  57. help='change case of keywords, CHOICE is one of {}'.format(
  58. ', '.join(f'"{x}"' for x in _CASE_CHOICES)))
  59. group.add_argument(
  60. '-i', '--identifiers',
  61. metavar='CHOICE',
  62. dest='identifier_case',
  63. choices=_CASE_CHOICES,
  64. help='change case of identifiers, CHOICE is one of {}'.format(
  65. ', '.join(f'"{x}"' for x in _CASE_CHOICES)))
  66. group.add_argument(
  67. '-l', '--language',
  68. metavar='LANG',
  69. dest='output_format',
  70. choices=['python', 'php'],
  71. help='output a snippet in programming language LANG, '
  72. 'choices are "python", "php"')
  73. group.add_argument(
  74. '--strip-comments',
  75. dest='strip_comments',
  76. action='store_true',
  77. default=False,
  78. help='remove comments')
  79. group.add_argument(
  80. '-r', '--reindent',
  81. dest='reindent',
  82. action='store_true',
  83. default=False,
  84. help='reindent statements')
  85. group.add_argument(
  86. '--indent_width',
  87. dest='indent_width',
  88. default=2,
  89. type=int,
  90. help='indentation width (defaults to 2 spaces)')
  91. group.add_argument(
  92. '--indent_after_first',
  93. dest='indent_after_first',
  94. action='store_true',
  95. default=False,
  96. help='indent after first line of statement (e.g. SELECT)')
  97. group.add_argument(
  98. '--indent_columns',
  99. dest='indent_columns',
  100. action='store_true',
  101. default=False,
  102. help='indent all columns by indent_width instead of keyword length')
  103. group.add_argument(
  104. '-a', '--reindent_aligned',
  105. action='store_true',
  106. default=False,
  107. help='reindent statements to aligned format')
  108. group.add_argument(
  109. '-s', '--use_space_around_operators',
  110. action='store_true',
  111. default=False,
  112. help='place spaces around mathematical operators')
  113. group.add_argument(
  114. '--wrap_after',
  115. dest='wrap_after',
  116. default=0,
  117. type=int,
  118. help='Column after which lists should be wrapped')
  119. group.add_argument(
  120. '--comma_first',
  121. dest='comma_first',
  122. default=False,
  123. type=bool,
  124. help='Insert linebreak before comma (default False)')
  125. group.add_argument(
  126. '--compact',
  127. dest='compact',
  128. default=False,
  129. type=bool,
  130. help='Try to produce more compact output (default False)')
  131. group.add_argument(
  132. '--encoding',
  133. dest='encoding',
  134. default='utf-8',
  135. help='Specify the input encoding (default utf-8)')
  136. return parser
  137. def _error(msg):
  138. """Print msg and optionally exit with return code exit_."""
  139. sys.stderr.write(f'[ERROR] {msg}\n')
  140. return 1
  141. def _process_file(filename, args):
  142. """Process a single file with the given formatting options.
  143. Returns 0 on success, 1 on error.
  144. """
  145. # Check for incompatible option combinations first
  146. if filename == '-' and args.inplace:
  147. return _error('Cannot use --in-place with stdin')
  148. # Read input
  149. if filename == '-': # read from stdin
  150. wrapper = TextIOWrapper(sys.stdin.buffer, encoding=args.encoding)
  151. try:
  152. data = wrapper.read()
  153. finally:
  154. wrapper.detach()
  155. else:
  156. try:
  157. with open(filename, encoding=args.encoding) as f:
  158. data = ''.join(f.readlines())
  159. except OSError as e:
  160. return _error(f'Failed to read {filename}: {e}')
  161. # Determine output destination
  162. close_stream = False
  163. if args.inplace:
  164. try:
  165. stream = open(filename, 'w', encoding=args.encoding)
  166. close_stream = True
  167. except OSError as e:
  168. return _error(f'Failed to open {filename}: {e}')
  169. elif args.outfile:
  170. try:
  171. stream = open(args.outfile, 'w', encoding=args.encoding)
  172. close_stream = True
  173. except OSError as e:
  174. return _error(f'Failed to open {args.outfile}: {e}')
  175. else:
  176. stream = sys.stdout
  177. # Format the SQL
  178. formatter_opts = vars(args)
  179. try:
  180. formatter_opts = sqlparse.formatter.validate_options(formatter_opts)
  181. except SQLParseError as e:
  182. return _error(f'Invalid options: {e}')
  183. s = sqlparse.format(data, **formatter_opts)
  184. stream.write(s)
  185. stream.flush()
  186. if close_stream:
  187. stream.close()
  188. return 0
  189. def main(args=None):
  190. parser = create_parser()
  191. args = parser.parse_args(args)
  192. # Validate argument combinations
  193. if len(args.filename) > 1:
  194. if args.outfile:
  195. return _error('Cannot use -o/--outfile with multiple files')
  196. if not args.inplace:
  197. return _error('Multiple files require --in-place flag')
  198. # Process all files
  199. exit_code = 0
  200. for filename in args.filename:
  201. result = _process_file(filename, args)
  202. if result != 0:
  203. exit_code = result
  204. # Continue processing remaining files even if one fails
  205. return exit_code