cli.py 5.6 KB

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