lib.py 3.4 KB

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