compat.py 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308
  1. from __future__ import absolute_import
  2. import errno
  3. import numbers
  4. import os
  5. import sys
  6. from .five import range, zip_longest
  7. if sys.platform == 'win32':
  8. try:
  9. import _winapi # noqa
  10. except ImportError: # pragma: no cover
  11. from _multiprocessing import win32 as _winapi # noqa
  12. else:
  13. _winapi = None # noqa
  14. try:
  15. import resource
  16. except ImportError: # pragma: no cover
  17. resource = None
  18. try:
  19. from io import UnsupportedOperation
  20. FILENO_ERRORS = (AttributeError, ValueError, UnsupportedOperation)
  21. except ImportError: # pragma: no cover
  22. # Py2
  23. FILENO_ERRORS = (AttributeError, ValueError) # noqa
  24. if sys.version_info > (2, 7, 5):
  25. buf_t, is_new_buffer = memoryview, True # noqa
  26. else:
  27. buf_t, is_new_buffer = buffer, False # noqa
  28. if hasattr(os, 'write'):
  29. __write__ = os.write
  30. if is_new_buffer:
  31. def send_offset(fd, buf, offset):
  32. return __write__(fd, buf[offset:])
  33. else: # Py<2.7.6
  34. def send_offset(fd, buf, offset): # noqa
  35. return __write__(fd, buf_t(buf, offset))
  36. else: # non-posix platform
  37. def send_offset(fd, buf, offset): # noqa
  38. raise NotImplementedError('send_offset')
  39. try:
  40. fsencode = os.fsencode
  41. fsdecode = os.fsdecode
  42. except AttributeError:
  43. def _fscodec():
  44. encoding = sys.getfilesystemencoding()
  45. if encoding == 'mbcs':
  46. errors = 'strict'
  47. else:
  48. errors = 'surrogateescape'
  49. def fsencode(filename):
  50. """
  51. Encode filename to the filesystem encoding with 'surrogateescape'
  52. error handler, return bytes unchanged. On Windows, use 'strict'
  53. error handler if the file system encoding is 'mbcs' (which is the
  54. default encoding).
  55. """
  56. if isinstance(filename, bytes):
  57. return filename
  58. elif isinstance(filename, str):
  59. return filename.encode(encoding, errors)
  60. else:
  61. raise TypeError("expect bytes or str, not %s"
  62. % type(filename).__name__)
  63. def fsdecode(filename):
  64. """
  65. Decode filename from the filesystem encoding with 'surrogateescape'
  66. error handler, return str unchanged. On Windows, use 'strict' error
  67. handler if the file system encoding is 'mbcs' (which is the default
  68. encoding).
  69. """
  70. if isinstance(filename, str):
  71. return filename
  72. elif isinstance(filename, bytes):
  73. return filename.decode(encoding, errors)
  74. else:
  75. raise TypeError("expect bytes or str, not %s"
  76. % type(filename).__name__)
  77. return fsencode, fsdecode
  78. fsencode, fsdecode = _fscodec()
  79. del _fscodec
  80. if sys.version_info[0] == 3:
  81. bytes = bytes
  82. else:
  83. _bytes = bytes
  84. # the 'bytes' alias in Python2 does not support an encoding argument.
  85. class bytes(_bytes): # noqa
  86. def __new__(cls, *args):
  87. if len(args) > 1:
  88. return _bytes(args[0]).encode(*args[1:])
  89. return _bytes(*args)
  90. def maybe_fileno(f):
  91. """Get object fileno, or :const:`None` if not defined."""
  92. if isinstance(f, numbers.Integral):
  93. return f
  94. try:
  95. return f.fileno()
  96. except FILENO_ERRORS:
  97. pass
  98. def get_fdmax(default=None):
  99. """Return the maximum number of open file descriptors
  100. on this system.
  101. :keyword default: Value returned if there's no file
  102. descriptor limit.
  103. """
  104. try:
  105. return os.sysconf('SC_OPEN_MAX')
  106. except:
  107. pass
  108. if resource is None: # Windows
  109. return default
  110. fdmax = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
  111. if fdmax == resource.RLIM_INFINITY:
  112. return default
  113. return fdmax
  114. def uniq(it):
  115. """Return all unique elements in ``it``, preserving order."""
  116. seen = set()
  117. return (seen.add(obj) or obj for obj in it if obj not in seen)
  118. try:
  119. closerange = os.closerange
  120. except AttributeError:
  121. def closerange(fd_low, fd_high): # noqa
  122. for fd in reversed(range(fd_low, fd_high)):
  123. try:
  124. os.close(fd)
  125. except OSError as exc:
  126. if exc.errno != errno.EBADF:
  127. raise
  128. def close_open_fds(keep=None):
  129. # must make sure this is 0-inclusive (Issue #celery/1882)
  130. keep = list(uniq(sorted(
  131. f for f in map(maybe_fileno, keep or []) if f is not None
  132. )))
  133. maxfd = get_fdmax(default=2048)
  134. kL, kH = iter([-1] + keep), iter(keep + [maxfd])
  135. for low, high in zip_longest(kL, kH):
  136. if low + 1 != high:
  137. closerange(low + 1, high)
  138. else:
  139. def close_open_fds(keep=None): # noqa
  140. keep = [maybe_fileno(f)
  141. for f in (keep or []) if maybe_fileno(f) is not None]
  142. for fd in reversed(range(get_fdmax(default=2048))):
  143. if fd not in keep:
  144. try:
  145. os.close(fd)
  146. except OSError as exc:
  147. if exc.errno != errno.EBADF:
  148. raise
  149. def get_errno(exc):
  150. """:exc:`socket.error` and :exc:`IOError` first got
  151. the ``.errno`` attribute in Py2.7"""
  152. try:
  153. return exc.errno
  154. except AttributeError:
  155. try:
  156. # e.args = (errno, reason)
  157. if isinstance(exc.args, tuple) and len(exc.args) == 2:
  158. return exc.args[0]
  159. except AttributeError:
  160. pass
  161. return 0
  162. try:
  163. import _posixsubprocess
  164. except ImportError:
  165. def spawnv_passfds(path, args, passfds):
  166. if sys.platform != 'win32':
  167. # when not using _posixsubprocess (on earlier python) and not on
  168. # windows, we want to keep stdout/stderr open...
  169. passfds = passfds + [
  170. maybe_fileno(sys.stdout),
  171. maybe_fileno(sys.stderr),
  172. ]
  173. pid = os.fork()
  174. if not pid:
  175. close_open_fds(keep=sorted(f for f in passfds if f))
  176. os.execv(fsencode(path), args)
  177. return pid
  178. else:
  179. def spawnv_passfds(path, args, passfds):
  180. passfds = sorted(passfds)
  181. errpipe_read, errpipe_write = os.pipe()
  182. try:
  183. return _posixsubprocess.fork_exec(
  184. args, [fsencode(path)], True, tuple(passfds), None, None,
  185. -1, -1, -1, -1, -1, -1, errpipe_read, errpipe_write,
  186. False, False, None)
  187. finally:
  188. os.close(errpipe_read)
  189. os.close(errpipe_write)
  190. if sys.platform == 'win32':
  191. def setblocking(handle, blocking):
  192. raise NotImplementedError('setblocking not implemented on win32')
  193. def isblocking(handle):
  194. raise NotImplementedError('isblocking not implemented on win32')
  195. else:
  196. from os import O_NONBLOCK
  197. from fcntl import fcntl, F_GETFL, F_SETFL
  198. def isblocking(handle): # noqa
  199. return not (fcntl(handle, F_GETFL) & O_NONBLOCK)
  200. def setblocking(handle, blocking): # noqa
  201. flags = fcntl(handle, F_GETFL, 0)
  202. fcntl(
  203. handle, F_SETFL,
  204. flags & (~O_NONBLOCK) if blocking else flags | O_NONBLOCK,
  205. )
  206. E_PSUTIL_MISSING = """
  207. On Windows, the ability to inspect memory usage requires the psutil library.
  208. You can install it using pip:
  209. $ pip install psutil
  210. """
  211. E_RESOURCE_MISSING = """
  212. Your platform ({0}) does not seem to have the `resource.getrusage' function.
  213. Please open an issue so that we can add support for this platform.
  214. """
  215. if sys.platform == 'win32':
  216. try:
  217. import psutil
  218. except ImportError: # pragma: no cover
  219. psutil = None # noqa
  220. def mem_rss():
  221. # type () -> int
  222. if psutil is None:
  223. raise ImportError(E_PSUTIL_MISSING.strip())
  224. return int(psutil.Process(os.getpid()).memory_info()[0] / 1024.0)
  225. else:
  226. try:
  227. from resource import getrusage, RUSAGE_SELF
  228. except ImportError: # pragma: no cover
  229. getrusage = RUSAGE_SELF = None # noqa
  230. if 'bsd' in sys.platform or sys.platform == 'darwin':
  231. # On BSD platforms :man:`getrusage(2)` ru_maxrss field is in bytes.
  232. def maxrss_to_kb(v):
  233. # type: (SupportsInt) -> int
  234. return int(v) / 1024.0
  235. else:
  236. # On Linux it's kilobytes.
  237. def maxrss_to_kb(v):
  238. # type: (SupportsInt) -> int
  239. return int(v)
  240. def mem_rss():
  241. # type () -> int
  242. if resource is None:
  243. raise ImportError(E_RESOURCE_MISSING.strip().format(sys.platform))
  244. return maxrss_to_kb(getrusage(RUSAGE_SELF).ru_maxrss)