disk.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136
  1. """
  2. Disk management utilities.
  3. """
  4. # Authors: Gael Varoquaux <gael dot varoquaux at normalesup dot org>
  5. # Lars Buitinck
  6. # Copyright (c) 2010 Gael Varoquaux
  7. # License: BSD Style, 3 clauses.
  8. import os
  9. import sys
  10. import time
  11. import errno
  12. import shutil
  13. from multiprocessing import util
  14. try:
  15. WindowsError
  16. except NameError:
  17. WindowsError = OSError
  18. def disk_used(path):
  19. """ Return the disk usage in a directory."""
  20. size = 0
  21. for file in os.listdir(path) + ['.']:
  22. stat = os.stat(os.path.join(path, file))
  23. if hasattr(stat, 'st_blocks'):
  24. size += stat.st_blocks * 512
  25. else:
  26. # on some platform st_blocks is not available (e.g., Windows)
  27. # approximate by rounding to next multiple of 512
  28. size += (stat.st_size // 512 + 1) * 512
  29. # We need to convert to int to avoid having longs on some systems (we
  30. # don't want longs to avoid problems we SQLite)
  31. return int(size / 1024.)
  32. def memstr_to_bytes(text):
  33. """ Convert a memory text to its value in bytes.
  34. """
  35. kilo = 1024
  36. units = dict(K=kilo, M=kilo ** 2, G=kilo ** 3)
  37. try:
  38. size = int(units[text[-1]] * float(text[:-1]))
  39. except (KeyError, ValueError) as e:
  40. raise ValueError(
  41. "Invalid literal for size give: %s (type %s) should be "
  42. "alike '10G', '500M', '50K'." % (text, type(text))) from e
  43. return size
  44. def mkdirp(d):
  45. """Ensure directory d exists (like mkdir -p on Unix)
  46. No guarantee that the directory is writable.
  47. """
  48. try:
  49. os.makedirs(d)
  50. except OSError as e:
  51. if e.errno != errno.EEXIST:
  52. raise
  53. # if a rmtree operation fails in rm_subdirs, wait for this much time (in secs),
  54. # then retry up to RM_SUBDIRS_N_RETRY times. If it still fails, raise the
  55. # exception. this mecanism ensures that the sub-process gc have the time to
  56. # collect and close the memmaps before we fail.
  57. RM_SUBDIRS_RETRY_TIME = 0.1
  58. RM_SUBDIRS_N_RETRY = 5
  59. def rm_subdirs(path, onerror=None):
  60. """Remove all subdirectories in this path.
  61. The directory indicated by `path` is left in place, and its subdirectories
  62. are erased.
  63. If onerror is set, it is called to handle the error with arguments (func,
  64. path, exc_info) where func is os.listdir, os.remove, or os.rmdir;
  65. path is the argument to that function that caused it to fail; and
  66. exc_info is a tuple returned by sys.exc_info(). If onerror is None,
  67. an exception is raised.
  68. """
  69. # NOTE this code is adapted from the one in shutil.rmtree, and is
  70. # just as fast
  71. names = []
  72. try:
  73. names = os.listdir(path)
  74. except os.error:
  75. if onerror is not None:
  76. onerror(os.listdir, path, sys.exc_info())
  77. else:
  78. raise
  79. for name in names:
  80. fullname = os.path.join(path, name)
  81. delete_folder(fullname, onerror=onerror)
  82. def delete_folder(folder_path, onerror=None, allow_non_empty=True):
  83. """Utility function to cleanup a temporary folder if it still exists."""
  84. if os.path.isdir(folder_path):
  85. if onerror is not None:
  86. shutil.rmtree(folder_path, False, onerror)
  87. else:
  88. # allow the rmtree to fail once, wait and re-try.
  89. # if the error is raised again, fail
  90. err_count = 0
  91. while True:
  92. files = os.listdir(folder_path)
  93. try:
  94. if len(files) == 0 or allow_non_empty:
  95. shutil.rmtree(
  96. folder_path, ignore_errors=False, onerror=None
  97. )
  98. util.debug(
  99. "Sucessfully deleted {}".format(folder_path))
  100. break
  101. else:
  102. raise OSError(
  103. "Expected empty folder {} but got {} "
  104. "files.".format(folder_path, len(files))
  105. )
  106. except (OSError, WindowsError):
  107. err_count += 1
  108. if err_count > RM_SUBDIRS_N_RETRY:
  109. # the folder cannot be deleted right now. It maybe
  110. # because some temporary files have not been deleted
  111. # yet.
  112. raise
  113. time.sleep(RM_SUBDIRS_RETRY_TIME)