cli.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. #!/usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. #
  4. # Copyright (C) 2016 Andi Albrecht, albrecht.andi@gmail.com
  5. #
  6. # This module is part of python-sqlparse and is released under
  7. # the BSD License: http://www.opensource.org/licenses/bsd-license.php
  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. import sqlparse
  22. from sqlparse.compat import PY2
  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 {0}'.format(
  51. ', '.join('"{0}"'.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 {0}'.format(
  58. ', '.join('"{0}"'.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. '-a', '--reindent_aligned',
  86. action='store_true',
  87. default=False,
  88. help='reindent statements to aligned format')
  89. group.add_argument(
  90. '-s', '--use_space_around_operators',
  91. action='store_true',
  92. default=False,
  93. help='place spaces around mathematical operators')
  94. group.add_argument(
  95. '--wrap_after',
  96. dest='wrap_after',
  97. default=0,
  98. type=int,
  99. help='Column after which lists should be wrapped')
  100. return parser
  101. def _error(msg):
  102. """Print msg and optionally exit with return code exit_."""
  103. sys.stderr.write('[ERROR] {0}\n'.format(msg))
  104. return 1
  105. def main(args=None):
  106. parser = create_parser()
  107. args = parser.parse_args(args)
  108. if args.filename == '-': # read from stdin
  109. data = sys.stdin.read()
  110. else:
  111. try:
  112. # TODO: Needs to deal with encoding
  113. data = ''.join(open(args.filename).readlines())
  114. except IOError as e:
  115. return _error('Failed to read {0}: {1}'.format(args.filename, e))
  116. if args.outfile:
  117. try:
  118. stream = open(args.outfile, 'w')
  119. except IOError as e:
  120. return _error('Failed to open {0}: {1}'.format(args.outfile, e))
  121. else:
  122. stream = sys.stdout
  123. formatter_opts = vars(args)
  124. try:
  125. formatter_opts = sqlparse.formatter.validate_options(formatter_opts)
  126. except SQLParseError as e:
  127. return _error('Invalid options: {0}'.format(e))
  128. s = sqlparse.format(data, **formatter_opts)
  129. if PY2:
  130. s = s.encode('utf-8', 'replace')
  131. stream.write(s)
  132. stream.flush()
  133. return 0