lib.py 2.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. import sys
  2. import locale
  3. import random
  4. import time
  5. import signal
  6. from contextlib import contextmanager
  7. @contextmanager
  8. def default_sigint():
  9. original_sigint_handler = signal.getsignal(signal.SIGINT)
  10. signal.signal(signal.SIGINT, signal.SIG_DFL)
  11. try:
  12. yield
  13. except:
  14. raise
  15. finally:
  16. signal.signal(signal.SIGINT, original_sigint_handler)
  17. def to_utf8(s):
  18. """Re-encode string from the default system encoding to UTF-8."""
  19. current = locale.getpreferredencoding()
  20. return s.decode(current).encode("UTF-8") if s and current != "UTF-8" else s
  21. def debug(obj, fd=sys.stderr):
  22. """Write obj to standard error."""
  23. string = str(obj.encode(get_encoding(fd), "backslashreplace")
  24. if isinstance(obj, unicode) else obj)
  25. fd.write(string + "\n")
  26. def catch_exceptions(exit_codes, fun, *args, **kwargs):
  27. """
  28. Catch exceptions on fun(*args, **kwargs) and return the exit code specified
  29. in the exit_codes dictionary. Return 0 if no exception is raised.
  30. """
  31. try:
  32. fun(*args, **kwargs)
  33. return 0
  34. except tuple(exit_codes.keys()) as exc:
  35. debug("[%s] %s" % (exc.__class__.__name__, exc))
  36. return exit_codes[exc.__class__]
  37. def get_encoding(fd):
  38. """Guess terminal encoding."""
  39. return fd.encoding or locale.getpreferredencoding()
  40. def first(it):
  41. """Return first element in iterable."""
  42. return it.next()
  43. def string_to_dict(string):
  44. """Return dictionary from string "key1=value1, key2=value2"."""
  45. if string:
  46. pairs = [s.strip() for s in string.split(",")]
  47. return dict(pair.split("=") for pair in pairs)
  48. def retriable_exceptions(fun, retriable_exceptions, max_retries=None):
  49. """Run function and retry on some exceptions (with exponential backoff)."""
  50. retry = 0
  51. while 1:
  52. try:
  53. return fun()
  54. except tuple(retriable_exceptions) as exc:
  55. retry += 1
  56. if type(exc) not in retriable_exceptions:
  57. raise exc
  58. elif max_retries is not None and retry > max_retries:
  59. debug("Retry limit reached, time to give up")
  60. raise exc
  61. else:
  62. seconds = random.uniform(0, 2**retry)
  63. debug("Retryable error {0}/{1}: {2}. Waiting {3} seconds".
  64. format(retry, max_retries or "-", type(exc).__name__, seconds))
  65. time.sleep(seconds)