lib.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. import os
  2. import sys
  3. import locale
  4. import random
  5. import time
  6. import signal
  7. from contextlib import contextmanager
  8. @contextmanager
  9. def default_sigint():
  10. original_sigint_handler = signal.getsignal(signal.SIGINT)
  11. signal.signal(signal.SIGINT, signal.SIG_DFL)
  12. try:
  13. yield
  14. finally:
  15. signal.signal(signal.SIGINT, original_sigint_handler)
  16. def to_utf8(s):
  17. """Re-encode string from the default system encoding to UTF-8."""
  18. current = locale.getpreferredencoding()
  19. if hasattr(s, 'decode'):#Python 3 workaround
  20. return s.decode(current).encode("UTF-8") if s and current != "UTF-8" else s
  21. else:
  22. if isinstance(s, bytes):
  23. s = bytes.decode(s)
  24. return s
  25. def debug(obj, fd=sys.stderr):
  26. """Write obj to standard error."""
  27. try:
  28. unicode
  29. should_encode = not isinstance(obj, unicode)
  30. except NameError:
  31. should_encode = False #Doing so is harmless in Python 3
  32. string = str(obj.encode(get_encoding(fd), "backslashreplace")
  33. if should_encode else obj)
  34. #Python 3 handling workaround
  35. if sys.version_info >= (3, 0) and isinstance(string, bytes):
  36. fd.buffer.write(string + "\n") #We write the encoding directly
  37. else:
  38. fd.write(string + "\n")
  39. def catch_exceptions(exit_codes, fun, *args, **kwargs):
  40. """
  41. Catch exceptions on fun(*args, **kwargs) and return the exit code specified
  42. in the exit_codes dictionary. Return 0 if no exception is raised.
  43. """
  44. try:
  45. fun(*args, **kwargs)
  46. return 0
  47. except tuple(exit_codes.keys()) as exc:
  48. debug("[{0}] {1}".format(exc.__class__.__name__, exc))
  49. return exit_codes[exc.__class__]
  50. def get_encoding(fd):
  51. """Guess terminal encoding."""
  52. return fd.encoding or locale.getpreferredencoding()
  53. def first(it):
  54. """Return first element in iterable."""
  55. return it.next()
  56. def string_to_dict(string):
  57. """Return dictionary from string "key1=value1, key2=value2"."""
  58. if string:
  59. pairs = [s.strip() for s in string.split(",")]
  60. return dict(pair.split("=") for pair in pairs)
  61. def get_first_existing_filename(prefixes, relative_path):
  62. """Get the first existing filename of relative_path seeking on prefixes directories."""
  63. for prefix in prefixes:
  64. path = os.path.join(prefix, relative_path)
  65. if os.path.exists(path):
  66. return path
  67. def retriable_exceptions(fun, retriable_exceptions, max_retries=None):
  68. """Run function and retry on some exceptions (with exponential backoff)."""
  69. retry = 0
  70. while 1:
  71. try:
  72. return fun()
  73. except tuple(retriable_exceptions) as exc:
  74. retry += 1
  75. if type(exc) not in retriable_exceptions:
  76. raise exc
  77. elif max_retries is not None and retry > max_retries:
  78. debug("[Retryable errors] Retry limit reached")
  79. raise exc
  80. else:
  81. seconds = random.uniform(0, 2**retry)
  82. message = ("[Retryable error {current_retry}/{total_retries}] " +
  83. "{error_type} ({error_msg}). Wait {wait_time} seconds").format(
  84. current_retry=retry,
  85. total_retries=max_retries or "-",
  86. error_type=type(exc).__name__,
  87. error_msg=str(exc) or "-",
  88. wait_time="%.1f" % seconds,
  89. )
  90. debug(message)
  91. time.sleep(seconds)