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.

105 lines
2.9 KiB

4 years ago
  1. # coding: utf-8
  2. #
  3. # Copyright © 2010—2014 Andrey Mikhaylenko and contributors
  4. #
  5. # This file is part of Argh.
  6. #
  7. # Argh is free software under terms of the GNU Lesser
  8. # General Public License version 3 (LGPLv3) as published by the Free
  9. # Software Foundation. See the file README.rst for copying conditions.
  10. #
  11. """
  12. Output Processing
  13. ~~~~~~~~~~~~~~~~~
  14. """
  15. import locale
  16. import sys
  17. from argh import compat
  18. __all__ = ['dump', 'encode_output', 'safe_input']
  19. def _input(prompt):
  20. # this function can be mocked up in tests
  21. if sys.version_info < (3,0):
  22. return raw_input(prompt)
  23. else:
  24. return input(prompt)
  25. def safe_input(prompt):
  26. """
  27. Prompts user for input. Correctly handles prompt message encoding.
  28. """
  29. if sys.version_info < (3,0):
  30. if isinstance(prompt, compat.text_type):
  31. # Python 2.x: unicode → bytes
  32. encoding = locale.getpreferredencoding() or 'utf-8'
  33. prompt = prompt.encode(encoding)
  34. else:
  35. if not isinstance(prompt, compat.text_type):
  36. # Python 3.x: bytes → unicode
  37. prompt = prompt.decode()
  38. return _input(prompt)
  39. def encode_output(value, output_file):
  40. """
  41. Encodes given value so it can be written to given file object.
  42. Value may be Unicode, binary string or any other data type.
  43. The exact behaviour depends on the Python version:
  44. Python 3.x
  45. `sys.stdout` is a `_io.TextIOWrapper` instance that accepts `str`
  46. (unicode) and breaks on `bytes`.
  47. It is OK to simply assume that everything is Unicode unless special
  48. handling is introduced in the client code.
  49. Thus, no additional processing is performed.
  50. Python 2.x
  51. `sys.stdout` is a file-like object that accepts `str` (bytes)
  52. and breaks when `unicode` is passed to `sys.stdout.write()`.
  53. We can expect both Unicode and bytes. They need to be encoded so as
  54. to match the file object encoding.
  55. The output is binary if the object doesn't explicitly require Unicode.
  56. """
  57. if sys.version_info > (3,0):
  58. # Python 3: whatever → unicode
  59. return compat.text_type(value)
  60. else:
  61. # Python 2: handle special cases
  62. stream_encoding = getattr(output_file, 'encoding', None)
  63. if stream_encoding:
  64. if stream_encoding.upper() == 'UTF-8':
  65. return compat.text_type(value)
  66. else:
  67. return value.encode(stream_encoding, 'ignore')
  68. else:
  69. # no explicit encoding requirements; force binary
  70. if isinstance(value, compat.text_type):
  71. # unicode → binary
  72. return value.encode('utf-8')
  73. else:
  74. return str(value)
  75. def dump(raw_data, output_file):
  76. """
  77. Writes given line to given output file.
  78. See :func:`encode_output` for details.
  79. """
  80. data = encode_output(raw_data, output_file)
  81. output_file.write(data)