diff --git a/nvitop/api/host.py b/nvitop/api/host.py index a45afce1..1a139434 100644 --- a/nvitop/api/host.py +++ b/nvitop/api/host.py @@ -23,17 +23,30 @@ from __future__ import annotations import os as _os -from typing import TYPE_CHECKING as _TYPE_CHECKING import psutil as _psutil from psutil import * # noqa: F403 # pylint: disable=wildcard-import,unused-wildcard-import,redefined-builtin - - -if _TYPE_CHECKING: - from collections.abc import Callable as _Callable - - -__all__ = [name for name in _psutil.__all__ if not name.startswith('_')] + [ +from psutil import ( # noqa: F401 + LINUX, + MACOS, + POSIX, + WINDOWS, + AccessDenied, + Error, + NoSuchProcess, + Process, + ZombieProcess, + boot_time, + cpu_percent, + pids, + swap_memory, + virtual_memory, +) +from psutil import Error as PsutilError # pylint: disable=reimported + + +__all__ = [ + 'PsutilError', 'getuser', 'hostname', 'load_average', @@ -45,24 +58,10 @@ 'WSL', 'WINDOWS_SUBSYSTEM_FOR_LINUX', ] -__all__[__all__.index('Error')] = 'PsutilError' +__all__ += [name for name in _psutil.__all__ if not name.startswith('_') and name != 'Error'] -PsutilError = Error = _psutil.Error # make alias -del Error - - -cpu_percent = _psutil.cpu_percent -virtual_memory = _psutil.virtual_memory -swap_memory = _psutil.swap_memory -Process = _psutil.Process -NoSuchProcess = _psutil.NoSuchProcess -ZombieProcess = _psutil.ZombieProcess -AccessDenied = _psutil.AccessDenied -POSIX = _psutil.POSIX -WINDOWS = _psutil.WINDOWS -LINUX = _psutil.LINUX -MACOS = _psutil.MACOS +del Error # renamed to PsutilError def getuser() -> str: @@ -99,7 +98,7 @@ def uptime() -> float: """Get the system uptime.""" import time as _time # pylint: disable=import-outside-toplevel - return _time.time() - _psutil.boot_time() + return _time.time() - boot_time() def memory_percent() -> float: @@ -112,19 +111,32 @@ def swap_percent() -> float: return swap_memory().percent -ppid_map: _Callable[[], dict[int, int]] = _psutil._ppid_map # pylint: disable=protected-access -"""Obtain a ``{pid: ppid, ...}`` dict for all running processes in one shot.""" +def ppid_map() -> dict[int, int]: + """Obtain a ``{pid: ppid, ...}`` dict for all running processes in one shot.""" + ret = {} + for pid in pids(): + try: + ret[pid] = Process(pid).ppid() + except (NoSuchProcess, ZombieProcess): # noqa: PERF203 + pass + return ret + + +try: + from psutil import _ppid_map as ppid_map # type: ignore[no-redef] # noqa: F811,RUF100 +except ImportError: + pass -def reverse_ppid_map() -> dict[int, list[int]]: # pylint: disable=function-redefined +def reverse_ppid_map() -> dict[int, list[int]]: """Obtain a ``{ppid: [pid, ...], ...}`` dict for all running processes in one shot.""" from collections import defaultdict # pylint: disable=import-outside-toplevel - tree = defaultdict(list) + ret = defaultdict(list) for pid, ppid in ppid_map().items(): - tree[ppid].append(pid) + ret[ppid].append(pid) - return tree + return ret if LINUX: