_common.py 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861
  1. # Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
  2. # Use of this source code is governed by a BSD-style license that can be
  3. # found in the LICENSE file.
  4. """Common objects shared by __init__.py and _ps*.py modules.
  5. Note: this module is imported by setup.py, so it should not import
  6. psutil or third-party modules.
  7. """
  8. import collections
  9. import enum
  10. import functools
  11. import os
  12. import socket
  13. import stat
  14. import sys
  15. import threading
  16. import warnings
  17. from socket import AF_INET
  18. from socket import SOCK_DGRAM
  19. from socket import SOCK_STREAM
  20. try:
  21. from socket import AF_INET6
  22. except ImportError:
  23. AF_INET6 = None
  24. try:
  25. from socket import AF_UNIX
  26. except ImportError:
  27. AF_UNIX = None
  28. PSUTIL_DEBUG = bool(os.getenv('PSUTIL_DEBUG'))
  29. _DEFAULT = object()
  30. # fmt: off
  31. __all__ = [
  32. # OS constants
  33. 'FREEBSD', 'BSD', 'LINUX', 'NETBSD', 'OPENBSD', 'MACOS', 'OSX', 'POSIX',
  34. 'SUNOS', 'WINDOWS',
  35. # connection constants
  36. 'CONN_CLOSE', 'CONN_CLOSE_WAIT', 'CONN_CLOSING', 'CONN_ESTABLISHED',
  37. 'CONN_FIN_WAIT1', 'CONN_FIN_WAIT2', 'CONN_LAST_ACK', 'CONN_LISTEN',
  38. 'CONN_NONE', 'CONN_SYN_RECV', 'CONN_SYN_SENT', 'CONN_TIME_WAIT',
  39. # net constants
  40. 'NIC_DUPLEX_FULL', 'NIC_DUPLEX_HALF', 'NIC_DUPLEX_UNKNOWN', # noqa: F822
  41. # process status constants
  42. 'STATUS_DEAD', 'STATUS_DISK_SLEEP', 'STATUS_IDLE', 'STATUS_LOCKED',
  43. 'STATUS_RUNNING', 'STATUS_SLEEPING', 'STATUS_STOPPED', 'STATUS_SUSPENDED',
  44. 'STATUS_TRACING_STOP', 'STATUS_WAITING', 'STATUS_WAKE_KILL',
  45. 'STATUS_WAKING', 'STATUS_ZOMBIE', 'STATUS_PARKED',
  46. # other constants
  47. 'ENCODING', 'ENCODING_ERRS', 'AF_INET6',
  48. # utility functions
  49. 'conn_tmap', 'deprecated_method', 'isfile_strict', 'memoize',
  50. 'parse_environ_block', 'path_exists_strict', 'usage_percent',
  51. 'supports_ipv6', 'sockfam_to_enum', 'socktype_to_enum', "wrap_numbers",
  52. 'open_text', 'open_binary', 'cat', 'bcat',
  53. 'bytes2human', 'conn_to_ntuple', 'debug',
  54. # shell utils
  55. 'hilite', 'term_supports_colors', 'print_color',
  56. ]
  57. # fmt: on
  58. # ===================================================================
  59. # --- OS constants
  60. # ===================================================================
  61. POSIX = os.name == "posix"
  62. WINDOWS = os.name == "nt"
  63. LINUX = sys.platform.startswith("linux")
  64. MACOS = sys.platform.startswith("darwin")
  65. OSX = MACOS # deprecated alias
  66. FREEBSD = sys.platform.startswith(("freebsd", "midnightbsd"))
  67. OPENBSD = sys.platform.startswith("openbsd")
  68. NETBSD = sys.platform.startswith("netbsd")
  69. BSD = FREEBSD or OPENBSD or NETBSD
  70. SUNOS = sys.platform.startswith(("sunos", "solaris"))
  71. AIX = sys.platform.startswith("aix")
  72. # ===================================================================
  73. # --- API constants
  74. # ===================================================================
  75. # Process.status()
  76. STATUS_RUNNING = "running"
  77. STATUS_SLEEPING = "sleeping"
  78. STATUS_DISK_SLEEP = "disk-sleep"
  79. STATUS_STOPPED = "stopped"
  80. STATUS_TRACING_STOP = "tracing-stop"
  81. STATUS_ZOMBIE = "zombie"
  82. STATUS_DEAD = "dead"
  83. STATUS_WAKE_KILL = "wake-kill"
  84. STATUS_WAKING = "waking"
  85. STATUS_IDLE = "idle" # Linux, macOS, FreeBSD
  86. STATUS_LOCKED = "locked" # FreeBSD
  87. STATUS_WAITING = "waiting" # FreeBSD
  88. STATUS_SUSPENDED = "suspended" # NetBSD
  89. STATUS_PARKED = "parked" # Linux
  90. # Process.net_connections() and psutil.net_connections()
  91. CONN_ESTABLISHED = "ESTABLISHED"
  92. CONN_SYN_SENT = "SYN_SENT"
  93. CONN_SYN_RECV = "SYN_RECV"
  94. CONN_FIN_WAIT1 = "FIN_WAIT1"
  95. CONN_FIN_WAIT2 = "FIN_WAIT2"
  96. CONN_TIME_WAIT = "TIME_WAIT"
  97. CONN_CLOSE = "CLOSE"
  98. CONN_CLOSE_WAIT = "CLOSE_WAIT"
  99. CONN_LAST_ACK = "LAST_ACK"
  100. CONN_LISTEN = "LISTEN"
  101. CONN_CLOSING = "CLOSING"
  102. CONN_NONE = "NONE"
  103. # net_if_stats()
  104. class NicDuplex(enum.IntEnum):
  105. NIC_DUPLEX_FULL = 2
  106. NIC_DUPLEX_HALF = 1
  107. NIC_DUPLEX_UNKNOWN = 0
  108. globals().update(NicDuplex.__members__)
  109. # sensors_battery()
  110. class BatteryTime(enum.IntEnum):
  111. POWER_TIME_UNKNOWN = -1
  112. POWER_TIME_UNLIMITED = -2
  113. globals().update(BatteryTime.__members__)
  114. # --- others
  115. ENCODING = sys.getfilesystemencoding()
  116. ENCODING_ERRS = sys.getfilesystemencodeerrors()
  117. # ===================================================================
  118. # --- Process.net_connections() 'kind' parameter mapping
  119. # ===================================================================
  120. conn_tmap = {
  121. "all": ([AF_INET, AF_INET6, AF_UNIX], [SOCK_STREAM, SOCK_DGRAM]),
  122. "tcp": ([AF_INET, AF_INET6], [SOCK_STREAM]),
  123. "tcp4": ([AF_INET], [SOCK_STREAM]),
  124. "udp": ([AF_INET, AF_INET6], [SOCK_DGRAM]),
  125. "udp4": ([AF_INET], [SOCK_DGRAM]),
  126. "inet": ([AF_INET, AF_INET6], [SOCK_STREAM, SOCK_DGRAM]),
  127. "inet4": ([AF_INET], [SOCK_STREAM, SOCK_DGRAM]),
  128. "inet6": ([AF_INET6], [SOCK_STREAM, SOCK_DGRAM]),
  129. }
  130. if AF_INET6 is not None:
  131. conn_tmap.update({
  132. "tcp6": ([AF_INET6], [SOCK_STREAM]),
  133. "udp6": ([AF_INET6], [SOCK_DGRAM]),
  134. })
  135. if AF_UNIX is not None and not SUNOS:
  136. conn_tmap.update({"unix": ([AF_UNIX], [SOCK_STREAM, SOCK_DGRAM])})
  137. # =====================================================================
  138. # --- Exceptions
  139. # =====================================================================
  140. class Error(Exception):
  141. """Base exception class. All other psutil exceptions inherit
  142. from this one.
  143. """
  144. __module__ = 'psutil'
  145. def _infodict(self, attrs):
  146. info = collections.OrderedDict()
  147. for name in attrs:
  148. value = getattr(self, name, None)
  149. if value or (name == "pid" and value == 0):
  150. info[name] = value
  151. return info
  152. def __str__(self):
  153. # invoked on `raise Error`
  154. info = self._infodict(("pid", "ppid", "name"))
  155. if info:
  156. details = "({})".format(
  157. ", ".join([f"{k}={v!r}" for k, v in info.items()])
  158. )
  159. else:
  160. details = None
  161. return " ".join([x for x in (getattr(self, "msg", ""), details) if x])
  162. def __repr__(self):
  163. # invoked on `repr(Error)`
  164. info = self._infodict(("pid", "ppid", "name", "seconds", "msg"))
  165. details = ", ".join([f"{k}={v!r}" for k, v in info.items()])
  166. return f"psutil.{self.__class__.__name__}({details})"
  167. class NoSuchProcess(Error):
  168. """Exception raised when a process with a certain PID doesn't
  169. or no longer exists.
  170. """
  171. __module__ = 'psutil'
  172. def __init__(self, pid, name=None, msg=None):
  173. Error.__init__(self)
  174. self.pid = pid
  175. self.name = name
  176. self.msg = msg or "process no longer exists"
  177. def __reduce__(self):
  178. return (self.__class__, (self.pid, self.name, self.msg))
  179. class ZombieProcess(NoSuchProcess):
  180. """Exception raised when querying a zombie process. This is
  181. raised on macOS, BSD and Solaris only, and not always: depending
  182. on the query the OS may be able to succeed anyway.
  183. On Linux all zombie processes are querable (hence this is never
  184. raised). Windows doesn't have zombie processes.
  185. """
  186. __module__ = 'psutil'
  187. def __init__(self, pid, name=None, ppid=None, msg=None):
  188. NoSuchProcess.__init__(self, pid, name, msg)
  189. self.ppid = ppid
  190. self.msg = msg or "PID still exists but it's a zombie"
  191. def __reduce__(self):
  192. return (self.__class__, (self.pid, self.name, self.ppid, self.msg))
  193. class AccessDenied(Error):
  194. """Exception raised when permission to perform an action is denied."""
  195. __module__ = 'psutil'
  196. def __init__(self, pid=None, name=None, msg=None):
  197. Error.__init__(self)
  198. self.pid = pid
  199. self.name = name
  200. self.msg = msg or ""
  201. def __reduce__(self):
  202. return (self.__class__, (self.pid, self.name, self.msg))
  203. class TimeoutExpired(Error):
  204. """Raised on Process.wait(timeout) if timeout expires and process
  205. is still alive.
  206. """
  207. __module__ = 'psutil'
  208. def __init__(self, seconds, pid=None, name=None):
  209. Error.__init__(self)
  210. self.seconds = seconds
  211. self.pid = pid
  212. self.name = name
  213. self.msg = f"timeout after {seconds} seconds"
  214. def __reduce__(self):
  215. return (self.__class__, (self.seconds, self.pid, self.name))
  216. # ===================================================================
  217. # --- utils
  218. # ===================================================================
  219. def usage_percent(used, total, round_=None):
  220. """Calculate percentage usage of 'used' against 'total'."""
  221. try:
  222. ret = (float(used) / total) * 100
  223. except ZeroDivisionError:
  224. return 0.0
  225. else:
  226. if round_ is not None:
  227. ret = round(ret, round_)
  228. return ret
  229. def memoize(fun):
  230. """A simple memoize decorator for functions supporting (hashable)
  231. positional arguments.
  232. It also provides a cache_clear() function for clearing the cache:
  233. >>> @memoize
  234. ... def foo()
  235. ... return 1
  236. ...
  237. >>> foo()
  238. 1
  239. >>> foo.cache_clear()
  240. >>>
  241. It supports:
  242. - functions
  243. - classes (acts as a @singleton)
  244. - staticmethods
  245. - classmethods
  246. It does NOT support:
  247. - methods
  248. """
  249. @functools.wraps(fun)
  250. def wrapper(*args, **kwargs):
  251. key = (args, frozenset(sorted(kwargs.items())))
  252. try:
  253. return cache[key]
  254. except KeyError:
  255. try:
  256. ret = cache[key] = fun(*args, **kwargs)
  257. except Exception as err:
  258. raise err from None
  259. return ret
  260. def cache_clear():
  261. """Clear cache."""
  262. cache.clear()
  263. cache = {}
  264. wrapper.cache_clear = cache_clear
  265. return wrapper
  266. def memoize_when_activated(fun):
  267. """A memoize decorator which is disabled by default. It can be
  268. activated and deactivated on request.
  269. For efficiency reasons it can be used only against class methods
  270. accepting no arguments.
  271. >>> class Foo:
  272. ... @memoize
  273. ... def foo()
  274. ... print(1)
  275. ...
  276. >>> f = Foo()
  277. >>> # deactivated (default)
  278. >>> foo()
  279. 1
  280. >>> foo()
  281. 1
  282. >>>
  283. >>> # activated
  284. >>> foo.cache_activate(self)
  285. >>> foo()
  286. 1
  287. >>> foo()
  288. >>> foo()
  289. >>>
  290. """
  291. @functools.wraps(fun)
  292. def wrapper(self):
  293. try:
  294. # case 1: we previously entered oneshot() ctx
  295. ret = self._cache[fun]
  296. except AttributeError:
  297. # case 2: we never entered oneshot() ctx
  298. try:
  299. return fun(self)
  300. except Exception as err:
  301. raise err from None
  302. except KeyError:
  303. # case 3: we entered oneshot() ctx but there's no cache
  304. # for this entry yet
  305. try:
  306. ret = fun(self)
  307. except Exception as err:
  308. raise err from None
  309. try:
  310. self._cache[fun] = ret
  311. except AttributeError:
  312. # multi-threading race condition, see:
  313. # https://github.com/giampaolo/psutil/issues/1948
  314. pass
  315. return ret
  316. def cache_activate(proc):
  317. """Activate cache. Expects a Process instance. Cache will be
  318. stored as a "_cache" instance attribute.
  319. """
  320. proc._cache = {}
  321. def cache_deactivate(proc):
  322. """Deactivate and clear cache."""
  323. try:
  324. del proc._cache
  325. except AttributeError:
  326. pass
  327. wrapper.cache_activate = cache_activate
  328. wrapper.cache_deactivate = cache_deactivate
  329. return wrapper
  330. def isfile_strict(path):
  331. """Same as os.path.isfile() but does not swallow EACCES / EPERM
  332. exceptions, see:
  333. http://mail.python.org/pipermail/python-dev/2012-June/120787.html.
  334. """
  335. try:
  336. st = os.stat(path)
  337. except PermissionError:
  338. raise
  339. except OSError:
  340. return False
  341. else:
  342. return stat.S_ISREG(st.st_mode)
  343. def path_exists_strict(path):
  344. """Same as os.path.exists() but does not swallow EACCES / EPERM
  345. exceptions. See:
  346. http://mail.python.org/pipermail/python-dev/2012-June/120787.html.
  347. """
  348. try:
  349. os.stat(path)
  350. except PermissionError:
  351. raise
  352. except OSError:
  353. return False
  354. else:
  355. return True
  356. def supports_ipv6():
  357. """Return True if IPv6 is supported on this platform."""
  358. if not socket.has_ipv6 or AF_INET6 is None:
  359. return False
  360. try:
  361. with socket.socket(AF_INET6, socket.SOCK_STREAM) as sock:
  362. sock.bind(("::1", 0))
  363. return True
  364. except OSError:
  365. return False
  366. def parse_environ_block(data):
  367. """Parse a C environ block of environment variables into a dictionary."""
  368. # The block is usually raw data from the target process. It might contain
  369. # trailing garbage and lines that do not look like assignments.
  370. ret = {}
  371. pos = 0
  372. # localize global variable to speed up access.
  373. WINDOWS_ = WINDOWS
  374. while True:
  375. next_pos = data.find("\0", pos)
  376. # nul byte at the beginning or double nul byte means finish
  377. if next_pos <= pos:
  378. break
  379. # there might not be an equals sign
  380. equal_pos = data.find("=", pos, next_pos)
  381. if equal_pos > pos:
  382. key = data[pos:equal_pos]
  383. value = data[equal_pos + 1 : next_pos]
  384. # Windows expects environment variables to be uppercase only
  385. if WINDOWS_:
  386. key = key.upper()
  387. ret[key] = value
  388. pos = next_pos + 1
  389. return ret
  390. def sockfam_to_enum(num):
  391. """Convert a numeric socket family value to an IntEnum member.
  392. If it's not a known member, return the numeric value itself.
  393. """
  394. try:
  395. return socket.AddressFamily(num)
  396. except ValueError:
  397. return num
  398. def socktype_to_enum(num):
  399. """Convert a numeric socket type value to an IntEnum member.
  400. If it's not a known member, return the numeric value itself.
  401. """
  402. try:
  403. return socket.SocketKind(num)
  404. except ValueError:
  405. return num
  406. def conn_to_ntuple(fd, fam, type_, laddr, raddr, status, status_map, pid=None):
  407. """Convert a raw connection tuple to a proper ntuple."""
  408. from . import _ntuples as ntp
  409. if fam in {socket.AF_INET, AF_INET6}:
  410. if laddr:
  411. laddr = ntp.addr(*laddr)
  412. if raddr:
  413. raddr = ntp.addr(*raddr)
  414. if type_ == socket.SOCK_STREAM and fam in {AF_INET, AF_INET6}:
  415. status = status_map.get(status, CONN_NONE)
  416. else:
  417. status = CONN_NONE # ignore whatever C returned to us
  418. fam = sockfam_to_enum(fam)
  419. type_ = socktype_to_enum(type_)
  420. if pid is None:
  421. return ntp.pconn(fd, fam, type_, laddr, raddr, status)
  422. else:
  423. return ntp.sconn(fd, fam, type_, laddr, raddr, status, pid)
  424. def broadcast_addr(addr):
  425. """Given the address ntuple returned by ``net_if_addrs()``
  426. calculates the broadcast address.
  427. """
  428. import ipaddress
  429. if not addr.address or not addr.netmask:
  430. return None
  431. if addr.family == socket.AF_INET:
  432. return str(
  433. ipaddress.IPv4Network(
  434. f"{addr.address}/{addr.netmask}", strict=False
  435. ).broadcast_address
  436. )
  437. if addr.family == socket.AF_INET6:
  438. return str(
  439. ipaddress.IPv6Network(
  440. f"{addr.address}/{addr.netmask}", strict=False
  441. ).broadcast_address
  442. )
  443. def deprecated_method(replacement):
  444. """A decorator which can be used to mark a method as deprecated
  445. 'replcement' is the method name which will be called instead.
  446. """
  447. def outer(fun):
  448. msg = (
  449. f"{fun.__name__}() is deprecated and will be removed; use"
  450. f" {replacement}() instead"
  451. )
  452. if fun.__doc__ is None:
  453. fun.__doc__ = msg
  454. @functools.wraps(fun)
  455. def inner(self, *args, **kwargs):
  456. warnings.warn(msg, category=DeprecationWarning, stacklevel=2)
  457. return getattr(self, replacement)(*args, **kwargs)
  458. return inner
  459. return outer
  460. class _WrapNumbers:
  461. """Watches numbers so that they don't overflow and wrap
  462. (reset to zero).
  463. """
  464. def __init__(self):
  465. self.lock = threading.Lock()
  466. self.cache = {}
  467. self.reminders = {}
  468. self.reminder_keys = {}
  469. def _add_dict(self, input_dict, name):
  470. assert name not in self.cache
  471. assert name not in self.reminders
  472. assert name not in self.reminder_keys
  473. self.cache[name] = input_dict
  474. self.reminders[name] = collections.defaultdict(int)
  475. self.reminder_keys[name] = collections.defaultdict(set)
  476. def _remove_dead_reminders(self, input_dict, name):
  477. """In case the number of keys changed between calls (e.g. a
  478. disk disappears) this removes the entry from self.reminders.
  479. """
  480. old_dict = self.cache[name]
  481. gone_keys = set(old_dict.keys()) - set(input_dict.keys())
  482. for gone_key in gone_keys:
  483. for remkey in self.reminder_keys[name][gone_key]:
  484. del self.reminders[name][remkey]
  485. del self.reminder_keys[name][gone_key]
  486. def run(self, input_dict, name):
  487. """Cache dict and sum numbers which overflow and wrap.
  488. Return an updated copy of `input_dict`.
  489. """
  490. if name not in self.cache:
  491. # This was the first call.
  492. self._add_dict(input_dict, name)
  493. return input_dict
  494. self._remove_dead_reminders(input_dict, name)
  495. old_dict = self.cache[name]
  496. new_dict = {}
  497. for key in input_dict:
  498. input_tuple = input_dict[key]
  499. try:
  500. old_tuple = old_dict[key]
  501. except KeyError:
  502. # The input dict has a new key (e.g. a new disk or NIC)
  503. # which didn't exist in the previous call.
  504. new_dict[key] = input_tuple
  505. continue
  506. bits = []
  507. for i in range(len(input_tuple)):
  508. input_value = input_tuple[i]
  509. old_value = old_tuple[i]
  510. remkey = (key, i)
  511. if input_value < old_value:
  512. # it wrapped!
  513. self.reminders[name][remkey] += old_value
  514. self.reminder_keys[name][key].add(remkey)
  515. bits.append(input_value + self.reminders[name][remkey])
  516. new_dict[key] = tuple(bits)
  517. self.cache[name] = input_dict
  518. return new_dict
  519. def cache_clear(self, name=None):
  520. """Clear the internal cache, optionally only for function 'name'."""
  521. with self.lock:
  522. if name is None:
  523. self.cache.clear()
  524. self.reminders.clear()
  525. self.reminder_keys.clear()
  526. else:
  527. self.cache.pop(name, None)
  528. self.reminders.pop(name, None)
  529. self.reminder_keys.pop(name, None)
  530. def cache_info(self):
  531. """Return internal cache dicts as a tuple of 3 elements."""
  532. with self.lock:
  533. return (self.cache, self.reminders, self.reminder_keys)
  534. def wrap_numbers(input_dict, name):
  535. """Given an `input_dict` and a function `name`, adjust the numbers
  536. which "wrap" (restart from zero) across different calls by adding
  537. "old value" to "new value" and return an updated dict.
  538. """
  539. with _wn.lock:
  540. return _wn.run(input_dict, name)
  541. _wn = _WrapNumbers()
  542. wrap_numbers.cache_clear = _wn.cache_clear
  543. wrap_numbers.cache_info = _wn.cache_info
  544. # The read buffer size for open() builtin. This (also) dictates how
  545. # much data we read(2) when iterating over file lines as in:
  546. # >>> with open(file) as f:
  547. # ... for line in f:
  548. # ... ...
  549. # Default per-line buffer size for binary files is 1K. For text files
  550. # is 8K. We use a bigger buffer (32K) in order to have more consistent
  551. # results when reading /proc pseudo files on Linux, see:
  552. # https://github.com/giampaolo/psutil/issues/2050
  553. # https://github.com/giampaolo/psutil/issues/708
  554. FILE_READ_BUFFER_SIZE = 32 * 1024
  555. def open_binary(fname):
  556. return open(fname, "rb", buffering=FILE_READ_BUFFER_SIZE)
  557. def open_text(fname):
  558. """Open a file in text mode by using the proper FS encoding and
  559. en/decoding error handlers.
  560. """
  561. # See:
  562. # https://github.com/giampaolo/psutil/issues/675
  563. # https://github.com/giampaolo/psutil/pull/733
  564. fobj = open( # noqa: SIM115
  565. fname,
  566. buffering=FILE_READ_BUFFER_SIZE,
  567. encoding=ENCODING,
  568. errors=ENCODING_ERRS,
  569. )
  570. try:
  571. # Dictates per-line read(2) buffer size. Defaults is 8k. See:
  572. # https://github.com/giampaolo/psutil/issues/2050#issuecomment-1013387546
  573. fobj._CHUNK_SIZE = FILE_READ_BUFFER_SIZE
  574. except AttributeError:
  575. pass
  576. except Exception:
  577. fobj.close()
  578. raise
  579. return fobj
  580. def cat(fname, fallback=_DEFAULT, _open=open_text):
  581. """Read entire file content and return it as a string. File is
  582. opened in text mode. If specified, `fallback` is the value
  583. returned in case of error, either if the file does not exist or
  584. it can't be read().
  585. """
  586. if fallback is _DEFAULT:
  587. with _open(fname) as f:
  588. return f.read()
  589. else:
  590. try:
  591. with _open(fname) as f:
  592. return f.read()
  593. except OSError:
  594. return fallback
  595. def bcat(fname, fallback=_DEFAULT):
  596. """Same as above but opens file in binary mode."""
  597. return cat(fname, fallback=fallback, _open=open_binary)
  598. def bytes2human(n, format="%(value).1f%(symbol)s"):
  599. """Used by various scripts. See: https://code.activestate.com/recipes/578019-bytes-to-human-human-to-bytes-converter/?in=user-4178764.
  600. >>> bytes2human(10000)
  601. '9.8K'
  602. >>> bytes2human(100001221)
  603. '95.4M'
  604. """
  605. symbols = ('B', 'K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y')
  606. prefix = {}
  607. for i, s in enumerate(symbols[1:]):
  608. prefix[s] = 1 << (i + 1) * 10
  609. for symbol in reversed(symbols[1:]):
  610. if abs(n) >= prefix[symbol]:
  611. value = float(n) / prefix[symbol]
  612. return format % locals()
  613. return format % dict(symbol=symbols[0], value=n)
  614. def get_procfs_path():
  615. """Return updated psutil.PROCFS_PATH constant."""
  616. return sys.modules['psutil'].PROCFS_PATH
  617. def decode(s):
  618. return s.decode(encoding=ENCODING, errors=ENCODING_ERRS)
  619. # =====================================================================
  620. # --- shell utils
  621. # =====================================================================
  622. @memoize
  623. def term_supports_colors(file=sys.stdout): # pragma: no cover
  624. if not hasattr(file, "isatty") or not file.isatty():
  625. return False
  626. try:
  627. file.fileno()
  628. except Exception: # noqa: BLE001
  629. return False
  630. return True
  631. def hilite(s, color=None, bold=False): # pragma: no cover
  632. """Return an highlighted version of 'string'."""
  633. if not term_supports_colors():
  634. return s
  635. attr = []
  636. colors = dict(
  637. blue='34',
  638. brown='33',
  639. darkgrey='30',
  640. green='32',
  641. grey='37',
  642. lightblue='36',
  643. red='91',
  644. violet='35',
  645. yellow='93',
  646. )
  647. colors[None] = '29'
  648. try:
  649. color = colors[color]
  650. except KeyError:
  651. msg = f"invalid color {color!r}; choose amongst {list(colors.keys())}"
  652. raise ValueError(msg) from None
  653. attr.append(color)
  654. if bold:
  655. attr.append('1')
  656. return f"\x1b[{';'.join(attr)}m{s}\x1b[0m"
  657. def print_color(
  658. s, color=None, bold=False, file=sys.stdout
  659. ): # pragma: no cover
  660. """Print a colorized version of string."""
  661. if not term_supports_colors():
  662. print(s, file=file)
  663. elif POSIX:
  664. print(hilite(s, color, bold), file=file)
  665. else:
  666. import ctypes
  667. DEFAULT_COLOR = 7
  668. GetStdHandle = ctypes.windll.Kernel32.GetStdHandle
  669. SetConsoleTextAttribute = (
  670. ctypes.windll.Kernel32.SetConsoleTextAttribute
  671. )
  672. colors = dict(green=2, red=4, brown=6, yellow=6)
  673. colors[None] = DEFAULT_COLOR
  674. try:
  675. color = colors[color]
  676. except KeyError:
  677. msg = (
  678. f"invalid color {color!r}; choose between"
  679. f" {list(colors.keys())!r}"
  680. )
  681. raise ValueError(msg) from None
  682. if bold and color <= 7:
  683. color += 8
  684. handle_id = -12 if file is sys.stderr else -11
  685. GetStdHandle.restype = ctypes.c_ulong
  686. handle = GetStdHandle(handle_id)
  687. SetConsoleTextAttribute(handle, color)
  688. try:
  689. print(s, file=file)
  690. finally:
  691. SetConsoleTextAttribute(handle, DEFAULT_COLOR)
  692. def debug(msg):
  693. """If PSUTIL_DEBUG env var is set, print a debug message to stderr."""
  694. if PSUTIL_DEBUG:
  695. import inspect
  696. fname, lineno, _, _lines, _index = inspect.getframeinfo(
  697. inspect.currentframe().f_back
  698. )
  699. if isinstance(msg, Exception):
  700. if isinstance(msg, OSError):
  701. # ...because str(exc) may contain info about the file name
  702. msg = f"ignoring {msg}"
  703. else:
  704. msg = f"ignoring {msg!r}"
  705. print( # noqa: T201
  706. f"psutil-debug [{fname}:{lineno}]> {msg}", file=sys.stderr
  707. )