__init__.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389
  1. # coding: utf-8
  2. import os
  3. import stat
  4. import sys
  5. import shutil
  6. import logging
  7. from six import reraise
  8. import library.python.func
  9. import library.python.strings
  10. logger = logging.getLogger(__name__)
  11. ERRORS = {
  12. 'SUCCESS': 0,
  13. 'PATH_NOT_FOUND': 3,
  14. 'ACCESS_DENIED': 5,
  15. 'SHARING_VIOLATION': 32,
  16. 'INSUFFICIENT_BUFFER': 122,
  17. 'DIR_NOT_EMPTY': 145,
  18. }
  19. RETRIABLE_FILE_ERRORS = (ERRORS['ACCESS_DENIED'], ERRORS['SHARING_VIOLATION'])
  20. RETRIABLE_DIR_ERRORS = (ERRORS['ACCESS_DENIED'], ERRORS['DIR_NOT_EMPTY'], ERRORS['SHARING_VIOLATION'])
  21. # Check if on Windows
  22. @library.python.func.lazy
  23. def on_win():
  24. return os.name == 'nt'
  25. class NotOnWindowsError(RuntimeError):
  26. def __init__(self, message):
  27. super(NotOnWindowsError, self).__init__(message)
  28. class DisabledOnWindowsError(RuntimeError):
  29. def __init__(self, message):
  30. super(DisabledOnWindowsError, self).__init__(message)
  31. class NoCTypesError(RuntimeError):
  32. def __init__(self, message):
  33. super(NoCTypesError, self).__init__(message)
  34. # Decorator for Windows-only functions
  35. def win_only(f):
  36. def f_wrapped(*args, **kwargs):
  37. if not on_win():
  38. raise NotOnWindowsError('Windows-only function is called, but platform is not Windows')
  39. return f(*args, **kwargs)
  40. return f_wrapped
  41. # Decorator for functions disabled on Windows
  42. def win_disabled(f):
  43. def f_wrapped(*args, **kwargs):
  44. if on_win():
  45. run_disabled()
  46. return f(*args, **kwargs)
  47. return f_wrapped
  48. def errorfix(f):
  49. if not on_win():
  50. return f
  51. def f_wrapped(*args, **kwargs):
  52. try:
  53. return f(*args, **kwargs)
  54. except WindowsError:
  55. tp, value, tb = sys.exc_info()
  56. fix_error(value)
  57. reraise(tp, value, tb)
  58. return f_wrapped
  59. # Decorator for diehard wrapper
  60. # On Windows platform retries to run function while specific WindowsError is thrown
  61. # On non-Windows platforms fallbacks to function itself
  62. def diehard(winerrors, tries=100, delay=1):
  63. def wrap(f):
  64. if not on_win():
  65. return f
  66. return lambda *args, **kwargs: run_diehard(f, winerrors, tries, delay, *args, **kwargs)
  67. return wrap
  68. if on_win():
  69. import msvcrt
  70. import time
  71. import library.python.strings
  72. _has_ctypes = True
  73. try:
  74. import ctypes
  75. from ctypes import wintypes
  76. except ImportError:
  77. _has_ctypes = False
  78. _INVALID_HANDLE_VALUE = -1
  79. _MOVEFILE_REPLACE_EXISTING = 0x1
  80. _MOVEFILE_WRITE_THROUGH = 0x8
  81. _SEM_FAILCRITICALERRORS = 0x1
  82. _SEM_NOGPFAULTERRORBOX = 0x2
  83. _SEM_NOALIGNMENTFAULTEXCEPT = 0x4
  84. _SEM_NOOPENFILEERRORBOX = 0x8
  85. _SYMBOLIC_LINK_FLAG_DIRECTORY = 0x1
  86. _CREATE_NO_WINDOW = 0x8000000
  87. _ATOMIC_RENAME_FILE_TRANSACTION_DEFAULT_TIMEOUT = 1000
  88. _HANDLE_FLAG_INHERIT = 0x1
  89. @win_only
  90. def require_ctypes(f):
  91. def f_wrapped(*args, **kwargs):
  92. if not _has_ctypes:
  93. raise NoCTypesError('No ctypes found')
  94. return f(*args, **kwargs)
  95. return f_wrapped
  96. # Run function in diehard mode (see diehard decorator commentary)
  97. @win_only
  98. def run_diehard(f, winerrors, tries, delay, *args, **kwargs):
  99. if isinstance(winerrors, int):
  100. winerrors = (winerrors,)
  101. ei = None
  102. for t in range(tries):
  103. if t:
  104. logger.debug('Diehard [errs %s]: try #%d in %s', ','.join(str(x) for x in winerrors), t, f)
  105. try:
  106. return f(*args, **kwargs)
  107. except WindowsError as e:
  108. if e.winerror not in winerrors:
  109. raise
  110. ei = sys.exc_info()
  111. time.sleep(delay)
  112. reraise(ei[0], ei[1], ei[2])
  113. # Placeholder for disabled functions
  114. @win_only
  115. def run_disabled(*args, **kwargs):
  116. raise DisabledOnWindowsError('Function called is disabled on Windows')
  117. class CustomWinError(WindowsError):
  118. def __init__(self, winerror, message='', filename=None):
  119. super(CustomWinError, self).__init__(winerror, message)
  120. self.message = message
  121. self.strerror = self.message if self.message else format_error(self.windows_error)
  122. self.filename = filename
  123. self.utf8 = True
  124. @win_only
  125. def unicode_path(path):
  126. return library.python.strings.to_unicode(path, library.python.strings.fs_encoding())
  127. @win_only
  128. @require_ctypes
  129. def format_error(error):
  130. if isinstance(error, WindowsError):
  131. error = error.winerror
  132. if not isinstance(error, int):
  133. return 'Unknown'
  134. return ctypes.FormatError(error)
  135. @win_only
  136. def fix_error(windows_error):
  137. if not windows_error.strerror:
  138. windows_error.strerror = format_error(windows_error)
  139. transcode_error(windows_error)
  140. @win_only
  141. def transcode_error(windows_error, to_enc='utf-8'):
  142. from_enc = 'utf-8' if getattr(windows_error, 'utf8', False) else library.python.strings.guess_default_encoding()
  143. if from_enc != to_enc:
  144. windows_error.strerror = library.python.strings.to_str(
  145. windows_error.strerror, to_enc=to_enc, from_enc=from_enc
  146. )
  147. setattr(windows_error, 'utf8', to_enc == 'utf-8')
  148. class Transaction(object):
  149. def __init__(self, timeout=None, description=''):
  150. self.timeout = timeout
  151. self.description = description
  152. @require_ctypes
  153. def __enter__(self):
  154. self._handle = ctypes.windll.ktmw32.CreateTransaction(None, 0, 0, 0, 0, self.timeout, self.description)
  155. if self._handle == _INVALID_HANDLE_VALUE:
  156. raise ctypes.WinError()
  157. return self._handle
  158. @require_ctypes
  159. def __exit__(self, t, v, tb):
  160. try:
  161. if not ctypes.windll.ktmw32.CommitTransaction(self._handle):
  162. raise ctypes.WinError()
  163. finally:
  164. ctypes.windll.kernel32.CloseHandle(self._handle)
  165. @win_only
  166. def file_handle(f):
  167. return msvcrt.get_osfhandle(f.fileno())
  168. # https://www.python.org/dev/peps/pep-0446/
  169. # http://mihalop.blogspot.ru/2014/05/python-subprocess-and-file-descriptors.html
  170. @require_ctypes
  171. @win_only
  172. def open_file(*args, **kwargs):
  173. f = open(*args, **kwargs)
  174. ctypes.windll.kernel32.SetHandleInformation(file_handle(f), _HANDLE_FLAG_INHERIT, 0)
  175. return f
  176. @win_only
  177. @require_ctypes
  178. def replace_file(src, dst):
  179. if not ctypes.windll.kernel32.MoveFileExW(
  180. unicode_path(src), unicode_path(dst), _MOVEFILE_REPLACE_EXISTING | _MOVEFILE_WRITE_THROUGH
  181. ):
  182. raise ctypes.WinError()
  183. @win_only
  184. @require_ctypes
  185. def replace_file_across_devices(src, dst):
  186. with Transaction(
  187. timeout=_ATOMIC_RENAME_FILE_TRANSACTION_DEFAULT_TIMEOUT,
  188. description='ya library.python.windows replace_file_across_devices',
  189. ) as transaction:
  190. if not ctypes.windll.kernel32.MoveFileTransactedW(
  191. unicode_path(src),
  192. unicode_path(dst),
  193. None,
  194. None,
  195. _MOVEFILE_REPLACE_EXISTING | _MOVEFILE_WRITE_THROUGH,
  196. transaction,
  197. ):
  198. raise ctypes.WinError()
  199. @win_only
  200. @require_ctypes
  201. def hardlink(src, lnk):
  202. if not ctypes.windll.kernel32.CreateHardLinkW(unicode_path(lnk), unicode_path(src), None):
  203. raise ctypes.WinError()
  204. # Requires SE_CREATE_SYMBOLIC_LINK_NAME privilege
  205. @win_only
  206. @win_disabled
  207. @require_ctypes
  208. def symlink_file(src, lnk):
  209. if not ctypes.windll.kernel32.CreateSymbolicLinkW(unicode_path(lnk), unicode_path(src), 0):
  210. raise ctypes.WinError()
  211. # Requires SE_CREATE_SYMBOLIC_LINK_NAME privilege
  212. @win_only
  213. @win_disabled
  214. @require_ctypes
  215. def symlink_dir(src, lnk):
  216. if not ctypes.windll.kernel32.CreateSymbolicLinkW(
  217. unicode_path(lnk), unicode_path(src), _SYMBOLIC_LINK_FLAG_DIRECTORY
  218. ):
  219. raise ctypes.WinError()
  220. @win_only
  221. @require_ctypes
  222. def lock_file(f, offset, length, raises=True):
  223. locked = ctypes.windll.kernel32.LockFile(
  224. file_handle(f), _low_dword(offset), _high_dword(offset), _low_dword(length), _high_dword(length)
  225. )
  226. if not raises:
  227. return bool(locked)
  228. if not locked:
  229. raise ctypes.WinError()
  230. @win_only
  231. @require_ctypes
  232. def unlock_file(f, offset, length, raises=True):
  233. unlocked = ctypes.windll.kernel32.UnlockFile(
  234. file_handle(f), _low_dword(offset), _high_dword(offset), _low_dword(length), _high_dword(length)
  235. )
  236. if not raises:
  237. return bool(unlocked)
  238. if not unlocked:
  239. raise ctypes.WinError()
  240. @win_only
  241. @require_ctypes
  242. def set_error_mode(mode):
  243. return ctypes.windll.kernel32.SetErrorMode(mode)
  244. @win_only
  245. def rmtree(path):
  246. def error_handler(func, handling_path, execinfo):
  247. e = execinfo[1]
  248. if e.winerror == ERRORS['PATH_NOT_FOUND']:
  249. handling_path = "\\\\?\\" + handling_path # handle path over 256 symbols
  250. if os.path.exists(path):
  251. return func(handling_path)
  252. if e.winerror == ERRORS['ACCESS_DENIED']:
  253. try:
  254. # removing of r/w directory with read-only files in it yields ACCESS_DENIED
  255. # which is not an insuperable obstacle https://bugs.python.org/issue19643
  256. os.chmod(handling_path, stat.S_IWRITE)
  257. except OSError:
  258. pass
  259. else:
  260. # propagate true last error if this attempt fails
  261. return func(handling_path)
  262. raise e
  263. shutil.rmtree(path, onerror=error_handler)
  264. # Don't display the Windows GPF dialog if the invoked program dies.
  265. # http://msdn.microsoft.com/en-us/library/windows/desktop/ms680621.aspx
  266. @win_only
  267. def disable_error_dialogs():
  268. set_error_mode(_SEM_NOGPFAULTERRORBOX | _SEM_FAILCRITICALERRORS)
  269. @win_only
  270. def default_process_creation_flags():
  271. return 0
  272. @require_ctypes
  273. def _low_dword(x):
  274. return ctypes.c_ulong(x & ((1 << 32) - 1))
  275. @require_ctypes
  276. def _high_dword(x):
  277. return ctypes.c_ulong((x >> 32) & ((1 << 32) - 1))
  278. @win_only
  279. @require_ctypes
  280. def get_current_process():
  281. handle = ctypes.windll.kernel32.GetCurrentProcess()
  282. if not handle:
  283. raise ctypes.WinError()
  284. return wintypes.HANDLE(handle)
  285. @win_only
  286. @require_ctypes
  287. def get_process_handle_count(proc_handle):
  288. assert isinstance(proc_handle, wintypes.HANDLE)
  289. GetProcessHandleCount = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.HANDLE, wintypes.POINTER(wintypes.DWORD))(
  290. ("GetProcessHandleCount", ctypes.windll.kernel32)
  291. )
  292. hndcnt = wintypes.DWORD()
  293. if not GetProcessHandleCount(proc_handle, ctypes.byref(hndcnt)):
  294. raise ctypes.WinError()
  295. return hndcnt.value
  296. @win_only
  297. @require_ctypes
  298. def set_handle_information(file, inherit=None, protect_from_close=None):
  299. for flag, value in [(inherit, 1), (protect_from_close, 2)]:
  300. if flag is not None:
  301. assert isinstance(flag, bool)
  302. if not ctypes.windll.kernel32.SetHandleInformation(
  303. file_handle(file), _low_dword(value), _low_dword(int(flag))
  304. ):
  305. raise ctypes.WinError()
  306. @win_only
  307. @require_ctypes
  308. def get_windows_directory():
  309. buf = ctypes.create_unicode_buffer(ctypes.wintypes.MAX_PATH)
  310. size = ctypes.windll.kernel32.GetWindowsDirectoryW(buf, ctypes.wintypes.MAX_PATH)
  311. if not size:
  312. raise ctypes.WinError()
  313. if size > ctypes.wintypes.MAX_PATH - 1:
  314. raise CustomWinError(ERRORS['INSUFFICIENT_BUFFER'])
  315. return ctypes.wstring_at(buf, size)