From 96a0fab34fee2387c0fbe46dd30df8129ec8d284 Mon Sep 17 00:00:00 2001 From: Xuehai Pan Date: Fri, 25 Apr 2025 05:50:41 +0800 Subject: [PATCH] lint(tui): add doctests and add type annotations in `nvitop.tui` (#164) --- .github/workflows/lint.yaml | 9 +- CHANGELOG.md | 1 + docs/source/spelling_wordlist.txt | 2 + nvitop/api/__init__.py | 2 +- nvitop/api/collector.py | 2 +- nvitop/api/device.py | 35 +- nvitop/api/host.py | 8 +- nvitop/api/libcuda.py | 4 +- nvitop/api/libcudart.py | 6 +- nvitop/api/libnvml.py | 10 +- nvitop/api/process.py | 4 +- nvitop/api/termcolor.py | 3 +- nvitop/api/utils.py | 86 +++- nvitop/callbacks/keras.py | 2 + nvitop/callbacks/lightning.py | 2 + nvitop/callbacks/pytorch_lightning.py | 2 + nvitop/callbacks/tensorboard.py | 2 + nvitop/cli.py | 14 +- nvitop/select.py | 11 +- nvitop/tui/__init__.py | 14 +- nvitop/tui/library/__init__.py | 59 ++- nvitop/tui/library/device.py | 126 +++--- nvitop/tui/library/displayable.py | 107 +++-- nvitop/tui/library/history.py | 206 +++++---- nvitop/tui/library/host.py | 56 +-- nvitop/tui/library/keybinding.py | 284 +++++++----- nvitop/tui/library/libcurses.py | 155 ++++--- nvitop/tui/library/messagebox.py | 345 +++++++------- nvitop/tui/library/mouse.py | 68 +-- nvitop/tui/library/process.py | 29 +- nvitop/tui/library/selection.py | 103 +++-- nvitop/tui/library/utils.py | 50 ++- nvitop/tui/library/widestring.py | 140 +++--- nvitop/tui/screens/__init__.py | 13 + nvitop/tui/screens/base.py | 32 ++ nvitop/tui/screens/environ.py | 145 +++--- nvitop/tui/screens/help.py | 36 +- nvitop/tui/screens/main/__init__.py | 203 ++++++--- nvitop/tui/screens/main/panels/__init__.py | 19 + nvitop/tui/screens/main/panels/base.py | 34 ++ .../tui/screens/main/{ => panels}/device.py | 149 +++--- nvitop/tui/screens/main/{ => panels}/host.py | 154 ++++--- .../tui/screens/main/{ => panels}/process.py | 423 ++++++++++-------- nvitop/tui/screens/metrics.py | 174 ++++--- nvitop/tui/screens/treeview.py | 271 ++++++----- nvitop/tui/tui.py | 196 ++++---- pyproject.toml | 10 - setup.py | 1 + 48 files changed, 2302 insertions(+), 1505 deletions(-) create mode 100644 nvitop/tui/screens/base.py create mode 100644 nvitop/tui/screens/main/panels/__init__.py create mode 100644 nvitop/tui/screens/main/panels/base.py rename nvitop/tui/screens/main/{ => panels}/device.py (84%) rename nvitop/tui/screens/main/{ => panels}/host.py (78%) rename nvitop/tui/screens/main/{ => panels}/process.py (69%) diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml index 1098844..bd36f28 100644 --- a/.github/workflows/lint.yaml +++ b/.github/workflows/lint.yaml @@ -27,7 +27,7 @@ jobs: id: py uses: actions/setup-python@v5 with: - python-version: "3.8 - 3.13" + python-version: "3.13" update-environment: true - name: Upgrade pip @@ -68,3 +68,10 @@ jobs: python -m pre_commit --version python -m pre_commit install --install-hooks python -m pre_commit run --all-files + + - name: xdoctest + run: | + python -m pip install xdoctest + python -m xdoctest --version + grep -P -l -r '^\s*(import|from) (doctest|unittest)\b' nvitop | + xargs -r -L 1 python -m xdoctest diff --git a/CHANGELOG.md b/CHANGELOG.md index 44610b9..078bf39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Show `%GMBW` in main screen by [@XuehaiPan](https://github.com/XuehaiPan) in [#156](https://github.com/XuehaiPan/nvitop/pull/156). +- Add doctests and add type annotations in `nvitop.tui` by [@XuehaiPan](https://github.com/XuehaiPan) in [#164](https://github.com/XuehaiPan/nvitop/pull/164). ### Changed diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt index d74fdf6..f30aa46 100644 --- a/docs/source/spelling_wordlist.txt +++ b/docs/source/spelling_wordlist.txt @@ -24,6 +24,7 @@ devicesnapshot displayables divmod docstring +doctest ecc enum env @@ -107,6 +108,7 @@ rtx runtime rw rx +selectable shader sm smi diff --git a/nvitop/api/__init__.py b/nvitop/api/__init__.py index 6866145..ba9479d 100644 --- a/nvitop/api/__init__.py +++ b/nvitop/api/__init__.py @@ -65,7 +65,7 @@ from nvitop.api.utils import ( # explicitly export these to appease mypy ) -__all__ = [ +__all__ = [ # noqa: RUF022 'NVMLError', 'nvmlCheckReturn', 'libnvml', diff --git a/nvitop/api/collector.py b/nvitop/api/collector.py index 0bd7b5a..cf60426 100644 --- a/nvitop/api/collector.py +++ b/nvitop/api/collector.py @@ -38,7 +38,7 @@ if TYPE_CHECKING: from collections.abc import Callable, Generator, Iterable -__all__ = ['take_snapshots', 'collect_in_background', 'ResourceMetricCollector'] +__all__ = ['ResourceMetricCollector', 'collect_in_background', 'take_snapshots'] class SnapshotResult(NamedTuple): # pylint: disable=missing-class-docstring diff --git a/nvitop/api/device.py b/nvitop/api/device.py index b5e8ae8..5752e22 100644 --- a/nvitop/api/device.py +++ b/nvitop/api/device.py @@ -115,7 +115,7 @@ import textwrap import threading import time from collections import OrderedDict -from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, overload +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple, overload from nvitop.api import libcuda, libcudart, libnvml from nvitop.api.process import GpuProcess @@ -132,20 +132,17 @@ from nvitop.api.utils import ( if TYPE_CHECKING: from collections.abc import Callable, Generator, Hashable, Iterable - from typing_extensions import ( - Literal, # Python 3.8+ - Self, # Python 3.11+ - ) + from typing_extensions import Self # Python 3.11+ __all__ = [ - 'Device', - 'PhysicalDevice', - 'MigDevice', 'CudaDevice', 'CudaMigDevice', - 'parse_cuda_visible_devices', + 'Device', + 'MigDevice', + 'PhysicalDevice', 'normalize_cuda_visible_devices', + 'parse_cuda_visible_devices', ] # Class definitions ################################################################################ @@ -266,7 +263,7 @@ class Device: # pylint: disable=too-many-instance-attributes,too-many-public-me # GPU UUID : `GPU-` # MIG UUID : `MIG-GPU-//` # MIG UUID (R470+): `MIG-` - UUID_PATTERN: re.Pattern = re.compile( + UUID_PATTERN: ClassVar[re.Pattern] = re.compile( r"""^ # full match (?:(?PMIG)-)? # prefix for MIG UUID (?:(?PGPU)-)? # prefix for GPU UUID @@ -283,8 +280,8 @@ class Device: # pylint: disable=too-many-instance-attributes,too-many-public-me flags=re.VERBOSE, ) - GPU_PROCESS_CLASS: type[GpuProcess] = GpuProcess - cuda: type[CudaDevice] = None # type: ignore[assignment] # defined in below + GPU_PROCESS_CLASS: ClassVar[type[GpuProcess]] = GpuProcess + cuda: ClassVar[type[CudaDevice]] = None # type: ignore[assignment] # defined in below """Shortcut for class :class:`CudaDevice`.""" _nvml_index: int | tuple[int, int] @@ -395,7 +392,7 @@ class Device: # pylint: disable=too-many-instance-attributes,too-many-public-me def from_indices( cls, indices: int | Iterable[int | tuple[int, int]] | None = None, - ) -> list[PhysicalDevice | MigDevice]: + ) -> list[Self]: """Return a list of devices of the given indices. Args: @@ -430,7 +427,7 @@ class Device: # pylint: disable=too-many-instance-attributes,too-many-public-me if isinstance(indices, int): indices = [indices] - return list(map(cls, indices)) # type: ignore[arg-type] + return list(map(cls, indices)) @staticmethod def from_cuda_visible_devices() -> list[CudaDevice]: @@ -2176,9 +2173,7 @@ class Device: # pylint: disable=too-many-instance-attributes,too-many-public-me ignore_function_not_found=True, ) # nvmlDeviceIsMigDeviceHandle returns c_uint - self._is_mig_device = bool( - is_mig_device, - ) + self._is_mig_device = bool(is_mig_device) return self._is_mig_device return False @@ -2679,8 +2674,7 @@ class MigDevice(Device): # pylint: disable=too-many-instance-attributes The attributes are defined in :attr:`SNAPSHOT_KEYS`. """ snapshot = super().as_snapshot() - snapshot.mig_index = self.mig_index # type: ignore[attr-defined] - + snapshot.mig_index = self.mig_index return snapshot SNAPSHOT_KEYS: ClassVar[list[str]] = [ @@ -2930,8 +2924,7 @@ class CudaDevice(Device): The attributes are defined in :attr:`SNAPSHOT_KEYS`. """ snapshot = super().as_snapshot() - snapshot.cuda_index = self.cuda_index # type: ignore[attr-defined] - + snapshot.cuda_index = self.cuda_index return snapshot diff --git a/nvitop/api/host.py b/nvitop/api/host.py index 1a13943..24d8061 100644 --- a/nvitop/api/host.py +++ b/nvitop/api/host.py @@ -46,17 +46,17 @@ from psutil import Error as PsutilError # pylint: disable=reimported __all__ = [ + 'WINDOWS_SUBSYSTEM_FOR_LINUX', + 'WSL', 'PsutilError', 'getuser', 'hostname', 'load_average', - 'uptime', 'memory_percent', - 'swap_percent', 'ppid_map', 'reverse_ppid_map', - 'WSL', - 'WINDOWS_SUBSYSTEM_FOR_LINUX', + 'swap_percent', + 'uptime', ] __all__ += [name for name in _psutil.__all__ if not name.startswith('_') and name != 'Error'] diff --git a/nvitop/api/libcuda.py b/nvitop/api/libcuda.py index 4713227..81550d2 100644 --- a/nvitop/api/libcuda.py +++ b/nvitop/api/libcuda.py @@ -522,7 +522,7 @@ def cuDriverGetVersion() -> str: """ fn = __cudaGetFunctionPointer('cuDriverGetVersion') - driver_version = _ctypes.c_int() + driver_version = _ctypes.c_int(0) ret = fn(_ctypes.byref(driver_version)) _cudaCheckReturn(ret) major = driver_version.value // 1000 @@ -733,7 +733,7 @@ def cuDeviceTotalMem(device: _c_CUdevice_t) -> int: """ fn = __cudaGetFunctionPointer('cuDeviceTotalMem') - bytes = _ctypes.c_size_t() # pylint: disable=redefined-builtin + bytes = _ctypes.c_size_t(0) # pylint: disable=redefined-builtin ret = fn(_ctypes.byref(bytes), device) _cudaCheckReturn(ret) return bytes.value diff --git a/nvitop/api/libcudart.py b/nvitop/api/libcudart.py index bff1e50..006d5d2 100644 --- a/nvitop/api/libcudart.py +++ b/nvitop/api/libcudart.py @@ -614,7 +614,7 @@ def cudaDriverGetVersion() -> str: """ fn = __cudaGetFunctionPointer('cudaDriverGetVersion') - driver_version = _ctypes.c_int() + driver_version = _ctypes.c_int(0) ret = fn(_ctypes.byref(driver_version)) _cudaCheckReturn(ret) major = driver_version.value // 1000 @@ -638,7 +638,7 @@ def cudaRuntimeGetVersion() -> str: """ fn = __cudaGetFunctionPointer('cudaRuntimeGetVersion') - runtime_version = _ctypes.c_int() + runtime_version = _ctypes.c_int(0) ret = fn(_ctypes.byref(runtime_version)) _cudaCheckReturn(ret) major = runtime_version.value // 1000 @@ -695,7 +695,7 @@ def cudaDeviceGetByPCIBusId(pciBusId: str) -> int: """ fn = __cudaGetFunctionPointer('cudaDeviceGetByPCIBusId') - device = _ctypes.c_int() + device = _ctypes.c_int(0) ret = fn(_ctypes.byref(device), _ctypes.c_char_p(pciBusId.encode('utf-8'))) _cudaCheckReturn(ret) return device.value diff --git a/nvitop/api/libnvml.py b/nvitop/api/libnvml.py index eca2f9a..02ade4d 100644 --- a/nvitop/api/libnvml.py +++ b/nvitop/api/libnvml.py @@ -55,13 +55,13 @@ __all__ = [ # will be updated in below 'NA', 'UINT_MAX', 'ULONGLONG_MAX', + 'NVMLError', 'nvmlCheckReturn', - 'nvmlQuery', - 'nvmlQueryFieldValues', 'nvmlInit', 'nvmlInitWithFlags', + 'nvmlQuery', + 'nvmlQueryFieldValues', 'nvmlShutdown', - 'NVMLError', ] @@ -592,7 +592,7 @@ if not _pynvml_installation_corrupted: 'usedGpuCcProtectedMemory': '%d B', } - __get_running_processes_version_suffix = None + __get_running_processes_version_suffix: str | None = None c_nvmlProcessInfo_t = c_nvmlProcessInfo_v3_t def __determine_get_running_processes_version_suffix() -> str: @@ -811,7 +811,7 @@ if not _pynvml_installation_corrupted: _fmt_: _ClassVar[dict[str, str]] = {'': '%d B'} nvmlMemory_v2 = getattr(_pynvml, 'nvmlMemory_v2', _ctypes.sizeof(c_nvmlMemory_v2_t) | 2 << 24) - __get_memory_info_version_suffix = None + __get_memory_info_version_suffix: str | None = None c_nvmlMemory_t = c_nvmlMemory_v2_t def __determine_get_memory_info_version_suffix() -> str: diff --git a/nvitop/api/process.py b/nvitop/api/process.py index fa69e9a..ef38168 100644 --- a/nvitop/api/process.py +++ b/nvitop/api/process.py @@ -49,7 +49,7 @@ if TYPE_CHECKING: from nvitop.api.device import Device -__all__ = ['HostProcess', 'GpuProcess', 'command_join'] +__all__ = ['GpuProcess', 'HostProcess', 'command_join'] if host.POSIX: @@ -237,7 +237,7 @@ class HostProcess(host.Process, ABC): def __repr__(self) -> str: """Return a string representation of the process.""" - return super().__repr__().replace(self.__class__.__module__ + '.', '', 1) + return super().__repr__().replace(f'{self.__class__.__module__}.', '', 1) def __reduce__(self) -> tuple[type[HostProcess], tuple[int]]: """Return state information for pickling.""" diff --git a/nvitop/api/termcolor.py b/nvitop/api/termcolor.py index 03a97eb..0f279eb 100644 --- a/nvitop/api/termcolor.py +++ b/nvitop/api/termcolor.py @@ -46,12 +46,11 @@ from __future__ import annotations import io import os import sys -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal if TYPE_CHECKING: from collections.abc import Iterable - from typing_extensions import Literal # Python 3.8+ Attribute = Literal[ 'bold', diff --git a/nvitop/api/utils.py b/nvitop/api/utils.py index 0bed20f..af27b9b 100644 --- a/nvitop/api/utils.py +++ b/nvitop/api/utils.py @@ -38,7 +38,7 @@ if TYPE_CHECKING: from collections.abc import Generator, Iterable, Iterator -__all__ = [ +__all__ = [ # noqa: RUF022 'NA', 'NaType', 'NotApplicable', @@ -96,8 +96,10 @@ def colored( bold, dark, underline, blink, reverse, concealed. Examples: - >>> colored('Hello, World!', 'red', 'on_grey', ['blue', 'blink']) - >>> colored('Hello, World!', 'green') + >>> colored('Hello, World!', 'red', 'on_grey', ['bold', 'blink']) # doctest: +ELLIPSIS + '...Hello, World!...' + >>> colored('Hello, World!', 'green') # doctest: +ELLIPSIS + '...Hello, World!...' """ if COLOR: return termcolor.colored(text, color=color, on_color=on_color, attrs=attrs) @@ -139,9 +141,10 @@ class NaType(str): def __new__(cls) -> NaType: """Get the singleton instance (:const:`nvitop.NA`).""" - if not hasattr(cls, '_instance'): - cls._instance = super().__new__(cls, 'N/A') - return cls._instance + instance = getattr(cls, '_instance', None) + if instance is None: + cls._instance = instance = super().__new__(cls, 'N/A') + return instance def __bool__(self) -> bool: """Convert :const:`NA` to :class:`bool` and return :data:`False`. @@ -208,9 +211,11 @@ class NaType(str): """Return :data:`math.nan` if the operand is a number (``NA - other``). >>> NA - 'str' + Traceback (most recent call last): + ... TypeError: unsupported operand type(s) for -: 'NaType' and 'str' >>> NA - NA - 'N/AN/A' + nan >>> NA + 1 nan >>> NA + 1.0 @@ -226,6 +231,8 @@ class NaType(str): """Return :data:`math.nan` if the operand is a number (``other - NA``). >>> 'str' - NA + Traceback (most recent call last): + ... TypeError: unsupported operand type(s) for -: 'str' and 'NaType' >>> 1 - NA nan @@ -274,9 +281,13 @@ class NaType(str): >>> NA / 1024.0 nan >>> NA / 0 - ZeroDivisionError: float division by zero + Traceback (most recent call last): + ... + ZeroDivisionError: ... >>> NA / 0.0 - ZeroDivisionError: float division by zero + Traceback (most recent call last): + ... + ZeroDivisionError: ... >>> NA / NA nan """ @@ -306,9 +317,13 @@ class NaType(str): >>> NA // 1024.0 nan >>> NA / 0 - ZeroDivisionError: float division by zero + Traceback (most recent call last): + ... + ZeroDivisionError: ... >>> NA / 0.0 - ZeroDivisionError: float division by zero + Traceback (most recent call last): + ... + ZeroDivisionError: ... >>> NA // NA nan """ @@ -338,9 +353,13 @@ class NaType(str): >>> NA % 1024.0 nan >>> NA % 0 - ZeroDivisionError: float modulo + Traceback (most recent call last): + ... + ZeroDivisionError: ... >>> NA % 0.0 - ZeroDivisionError: float modulo + Traceback (most recent call last): + ... + ZeroDivisionError: ... """ if isinstance(other, (int, float)): return float(self) % other @@ -368,9 +387,13 @@ class NaType(str): >>> divmod(NA, 1024.0) (nan, nan) >>> divmod(NA, 0) - ZeroDivisionError: float floor division by zero + Traceback (most recent call last): + ... + ZeroDivisionError: ... >>> divmod(NA, 0.0) - ZeroDivisionError: float floor division by zero + Traceback (most recent call last): + ... + ZeroDivisionError: ... """ return (self // other, self % other) @@ -409,7 +432,7 @@ class NaType(str): return abs(float(self)) def __round__(self, ndigits: int | None = None) -> int | float: - """Round :const:`nvitop.NA` to ``ndigits`` decimal places, defaulting to :const:`0`. + """Round :const:`nvitop.NA` to ``ndigits`` decimal places, defaulting to :data:`None`. If ``ndigits`` is omitted or :data:`None`, returns :const:`0`, otherwise returns :data:`math.nan`. @@ -504,7 +527,7 @@ SIZE_UNITS: dict[str | None, int] = { } """Units of storage and memory measurements.""" SIZE_PATTERN: re.Pattern = re.compile( - r'^\s*\+?\s*(?P\d+(?:\.\d+)?)\s*(?P[KMGTP]i?B?|B?)\s*$', + r'^\s*\+?\s*(?P\d+(?:\.\d+)?)\s*(?P([KMGTP]i?)?)B?\s*$', flags=re.IGNORECASE, ) """The regex pattern for human readable size.""" @@ -558,6 +581,8 @@ def human2bytes(s: int | str, /) -> int: If cannot convert the given size string. Examples: + >>> human2bytes('200') + 200 >>> human2bytes('500B') 500 >>> human2bytes('10k') @@ -576,12 +601,12 @@ def human2bytes(s: int | str, /) -> int: return s raise ValueError(f'Cannot convert {s!r} to bytes.') - match = SIZE_PATTERN.match(s) + match = SIZE_PATTERN.fullmatch(s) if match is None: raise ValueError(f'Cannot convert {s!r} to bytes.') - size, unit = match.groups() - unit = unit.upper().replace('I', 'i').replace('B', '') + 'B' - return int(float(size) * SIZE_UNITS[unit]) + size, unit = match.group('size', 'unit') + unit = unit.upper().replace('I', 'i') + return int(float(size) * SIZE_UNITS[f'{unit}B']) def timedelta2human( @@ -638,8 +663,8 @@ class Snapshot: def __init__(self, real: Any, **items: Any) -> None: """Initialize a new :class:`Snapshot` object with the given attributes.""" - self.real = real - self.timestamp = time.time() + object.__setattr__(self, 'real', real) + object.__setattr__(self, 'timestamp', time.time()) for key, value in items.items(): setattr(self, key, value) @@ -690,6 +715,15 @@ class Snapshot: """Support ``snapshot['name'] = value`` syntax.""" setattr(self, name, value) + def __setattr__(self, name: str, value: Any) -> None: + """Set or update a member of the instance. + + If the attribute is not defined, set it to the snapshot object. + """ + if name in ('real', 'timestamp'): + raise AttributeError(f'Cannot set attribute {name!r} of {self.__class__.__name__!r}') + super().__setattr__(name, value) + def __iter__(self) -> Iterator[str]: """Support ``for name in snapshot`` syntax and ``*`` tuple unpack ``[*snapshot]`` syntax.""" @@ -757,3 +791,9 @@ def memoize_when_activated(method: Method, /) -> Method: wrapped.cache_activate = cache_activate # type: ignore[attr-defined] wrapped.cache_deactivate = cache_deactivate # type: ignore[attr-defined] return wrapped # type: ignore[return-value] + + +if __name__ == '__main__': + import doctest + + doctest.testmod() diff --git a/nvitop/callbacks/keras.py b/nvitop/callbacks/keras.py index 0754102..f520689 100644 --- a/nvitop/callbacks/keras.py +++ b/nvitop/callbacks/keras.py @@ -18,6 +18,8 @@ # pylint: disable=missing-module-docstring,missing-function-docstring # pylint: disable=unused-argument,attribute-defined-outside-init +# mypy: ignore-errors + from __future__ import annotations import re diff --git a/nvitop/callbacks/lightning.py b/nvitop/callbacks/lightning.py index f7425c3..f43f4af 100644 --- a/nvitop/callbacks/lightning.py +++ b/nvitop/callbacks/lightning.py @@ -18,6 +18,8 @@ # pylint: disable=missing-module-docstring,missing-function-docstring # pylint: disable=unused-argument,attribute-defined-outside-init +# mypy: ignore-errors + from __future__ import annotations import time diff --git a/nvitop/callbacks/pytorch_lightning.py b/nvitop/callbacks/pytorch_lightning.py index f9853b4..f015f0b 100644 --- a/nvitop/callbacks/pytorch_lightning.py +++ b/nvitop/callbacks/pytorch_lightning.py @@ -18,6 +18,8 @@ # pylint: disable=missing-module-docstring,missing-function-docstring # pylint: disable=unused-argument,attribute-defined-outside-init +# mypy: ignore-errors + from __future__ import annotations import time diff --git a/nvitop/callbacks/tensorboard.py b/nvitop/callbacks/tensorboard.py index 61a8a17..b835366 100644 --- a/nvitop/callbacks/tensorboard.py +++ b/nvitop/callbacks/tensorboard.py @@ -17,6 +17,8 @@ # pylint: disable=missing-module-docstring +# mypy: ignore-errors + from __future__ import annotations import warnings diff --git a/nvitop/cli.py b/nvitop/cli.py index 6ab7d97..236b826 100644 --- a/nvitop/cli.py +++ b/nvitop/cli.py @@ -95,10 +95,10 @@ def parse_arguments() -> argparse.Namespace: help='Process status update interval in seconds. (default: 2)', ) parser.add_argument( - '--ascii', '--no-unicode', + '--ascii', '-U', - dest='ascii', + dest='no_unicode', action='store_true', help='Use ASCII characters only, which is useful for terminals without Unicode support.', ) @@ -294,7 +294,7 @@ def main() -> int: args.monitor = mode if not setlocale_utf8(): - args.ascii = True + args.no_unicode = True try: device_count = Device.count() @@ -352,7 +352,7 @@ def main() -> int: tui = TUI( devices, filters, - ascii=args.ascii, + no_unicode=args.no_unicode, mode=args.monitor, interval=args.interval, win=win, @@ -364,7 +364,7 @@ def main() -> int: messages.append(f'ERROR: Failed to initialize `curses` ({ex})') if tui is None: - tui = TUI(devices, filters, ascii=args.ascii) + tui = TUI(devices, filters, no_unicode=args.no_unicode) if not sys.stdout.isatty(): parent = HostProcess().parent() if parent is not None: @@ -416,8 +416,8 @@ def main() -> int: pip3 install --upgrade pipx pipx run nvitop """, - ) - messages.append(message.strip() + '\n') + ).strip() + messages.append(f'{message}\n') if len(messages) > 0: for message in messages: diff --git a/nvitop/select.py b/nvitop/select.py index 5091f8a..0ef366f 100644 --- a/nvitop/select.py +++ b/nvitop/select.py @@ -62,7 +62,7 @@ import math import os import sys import warnings -from typing import TYPE_CHECKING, overload +from typing import TYPE_CHECKING, Literal, overload from nvitop.api import Device, GpuProcess, Snapshot, colored, host, human2bytes, libnvml from nvitop.version import __version__ @@ -70,7 +70,6 @@ from nvitop.version import __version__ if TYPE_CHECKING: from collections.abc import Callable, Iterable, Sequence - from typing_extensions import Literal # Python 3.8+ __all__ = ['select_devices'] @@ -220,7 +219,7 @@ def select_devices( for device in devices: available_devices.extend(dev.as_snapshot() for dev in device.to_leaf_devices()) for device in available_devices: - device.loosen_constraints = 0 # type: ignore[attr-defined] + device.loosen_constraints = 0 if len(free_accounts) > 0: with GpuProcess.failsafe(): @@ -229,15 +228,15 @@ def select_devices( for process in device.real.processes().values(): if process.username() in free_accounts: as_free_memory += process.gpu_memory() - device.memory_free += as_free_memory # type: ignore[attr-defined] - device.memory_used -= as_free_memory # type: ignore[attr-defined] + device.memory_free += as_free_memory + device.memory_used -= as_free_memory def filter_func( criteria: Callable[[Snapshot], bool], original_criteria: Callable[[Snapshot], bool], ) -> Callable[[Snapshot], bool]: def wrapped(device: Snapshot) -> bool: - device.loosen_constraints += int(not original_criteria(device)) # type: ignore[attr-defined] + device.loosen_constraints += int(not original_criteria(device)) return criteria(device) return wrapped diff --git a/nvitop/tui/__init__.py b/nvitop/tui/__init__.py index 173c38d..cd8ba68 100644 --- a/nvitop/tui/__init__.py +++ b/nvitop/tui/__init__.py @@ -4,7 +4,7 @@ # pylint: disable=missing-module-docstring from nvitop.tui.library import ( - SUPERUSER, + IS_SUPERUSER, USERNAME, Device, colored, @@ -13,3 +13,15 @@ from nvitop.tui.library import ( setlocale_utf8, ) from nvitop.tui.tui import TUI + + +__all__ = [ + 'IS_SUPERUSER', + 'TUI', + 'USERNAME', + 'Device', + 'colored', + 'libcurses', + 'set_color', + 'setlocale_utf8', +] diff --git a/nvitop/tui/library/__init__.py b/nvitop/tui/library/__init__.py index 5566794..6d0d2a4 100644 --- a/nvitop/tui/library/__init__.py +++ b/nvitop/tui/library/__init__.py @@ -4,7 +4,7 @@ # pylint: disable=missing-module-docstring from nvitop.tui.library import host -from nvitop.tui.library.device import Device +from nvitop.tui.library.device import Device, MigDevice from nvitop.tui.library.displayable import Displayable, DisplayableContainer from nvitop.tui.library.history import BufferedHistoryGraph, HistoryGraph from nvitop.tui.library.keybinding import ( @@ -18,20 +18,22 @@ from nvitop.tui.library.keybinding import ( normalize_keybinding, ) from nvitop.tui.library.libcurses import libcurses, setlocale_utf8 -from nvitop.tui.library.messagebox import MessageBox, send_signal +from nvitop.tui.library.messagebox import MessageBox from nvitop.tui.library.mouse import MouseEvent from nvitop.tui.library.process import GpuProcess, HostProcess from nvitop.tui.library.selection import Selection from nvitop.tui.library.utils import ( HOSTNAME, + IS_SUPERUSER, + IS_WINDOWS, + IS_WINDOWS_SUBSYSTEM_FOR_LINUX, + IS_WSL, LARGE_INTEGER, NA, - SUPERUSER, USER_CONTEXT, USERNAME, - WINDOWS, - WSL, GiB, + NaType, Snapshot, bytes2human, colored, @@ -42,3 +44,50 @@ from nvitop.tui.library.utils import ( ttl_cache, ) from nvitop.tui.library.widestring import WideString, wcslen + + +__all__ = [ + 'ALT_KEY', + 'ANYKEY', + 'HOSTNAME', + 'IS_SUPERUSER', + 'IS_WINDOWS', + 'IS_WINDOWS_SUBSYSTEM_FOR_LINUX', + 'IS_WSL', + 'LARGE_INTEGER', + 'NA', + 'PASSIVE_ACTION', + 'QUANT_KEY', + 'SPECIAL_KEYS', + 'USERNAME', + 'USER_CONTEXT', + 'BufferedHistoryGraph', + 'Device', + 'Displayable', + 'DisplayableContainer', + 'GiB', + 'GpuProcess', + 'HistoryGraph', + 'HostProcess', + 'KeyBuffer', + 'KeyMaps', + 'MessageBox', + 'MigDevice', + 'MouseEvent', + 'NaType', + 'Selection', + 'Snapshot', + 'WideString', + 'bytes2human', + 'colored', + 'cut_string', + 'host', + 'libcurses', + 'make_bar', + 'normalize_keybinding', + 'set_color', + 'setlocale_utf8', + 'timedelta2human', + 'ttl_cache', + 'wcslen', +] diff --git a/nvitop/tui/library/device.py b/nvitop/tui/library/device.py index 2b04fdd..8973384 100644 --- a/nvitop/tui/library/device.py +++ b/nvitop/tui/library/device.py @@ -3,23 +3,40 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring -from nvitop.api import NA, libnvml, ttl_cache, utilization2string +from __future__ import annotations + +import enum +from typing import Any, ClassVar, Literal + +from nvitop.api import NA, Snapshot, libnvml, ttl_cache, utilization2string from nvitop.api import MigDevice as MigDeviceBase from nvitop.api import PhysicalDevice as DeviceBase -from nvitop.tui.library.process import GpuProcess +from nvitop.tui.library.process import GpuProcess, GpuProcessBase -__all__ = ['Device', 'NA'] +__all__ = ['Device', 'MigDevice'] + + +class LoadingIntensity(enum.IntEnum): + LIGHT = 0 + MODERATE = 1 + HEAVY = 2 + + def color(self) -> str: + if self == LoadingIntensity.LIGHT: + return 'green' + if self == LoadingIntensity.MODERATE: + return 'yellow' + return 'red' class Device(DeviceBase): - GPU_PROCESS_CLASS = GpuProcess + GPU_PROCESS_CLASS: ClassVar[type[GpuProcessBase]] = GpuProcess - MEMORY_UTILIZATION_THRESHOLDS = (10, 80) - GPU_UTILIZATION_THRESHOLDS = (10, 75) - INTENSITY2COLOR = {'light': 'green', 'moderate': 'yellow', 'heavy': 'red'} + MEMORY_UTILIZATION_THRESHOLDS: ClassVar[tuple[int, int]] = (10, 80) + GPU_UTILIZATION_THRESHOLDS: ClassVar[tuple[int, int]] = (10, 75) - SNAPSHOT_KEYS = [ + SNAPSHOT_KEYS: ClassVar[list[str]] = [ 'name', 'bus_id', 'memory_used', @@ -58,26 +75,28 @@ class Device(DeviceBase): 'display_color', ] - def __init__(self, *args, **kwargs): + def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - self._snapshot = None - self.tuple_index = (self.index,) if isinstance(self.index, int) else self.index - self.display_index = ':'.join(map(str, self.tuple_index)) + self._snapshot: Snapshot | None = None + self.tuple_index: tuple[int] | tuple[int, int] = ( + (self.index,) if isinstance(self.index, int) else self.index + ) + self.display_index: str = ':'.join(map(str, self.tuple_index)) - def as_snapshot(self): + def as_snapshot(self) -> Snapshot: self._snapshot = super().as_snapshot() self._snapshot.tuple_index = self.tuple_index self._snapshot.display_index = self.display_index return self._snapshot @property - def snapshot(self): + def snapshot(self) -> Snapshot: if self._snapshot is None: - self.as_snapshot() + self._snapshot = self.as_snapshot() return self._snapshot - def mig_devices(self): + def mig_devices(self) -> list[MigDevice]: # type: ignore[override] mig_devices = [] if self.is_mig_mode_enabled(): @@ -105,89 +124,90 @@ class Device(DeviceBase): compute_mode = ttl_cache(ttl=5.0)(DeviceBase.compute_mode) mig_mode = ttl_cache(ttl=5.0)(DeviceBase.mig_mode) - def memory_percent_string(self): # in percentage + def memory_percent_string(self) -> str: # in percentage return utilization2string(self.memory_percent()) - def memory_utilization_string(self): # in percentage + def memory_utilization_string(self) -> str: # in percentage return utilization2string(self.memory_utilization()) - def gpu_utilization_string(self): # in percentage + def gpu_utilization_string(self) -> str: # in percentage return utilization2string(self.gpu_utilization()) - def fan_speed_string(self): # in percentage + def fan_speed_string(self) -> str: # in percentage return utilization2string(self.fan_speed()) - def temperature_string(self): # in Celsius + def temperature_string(self) -> str: # in Celsius temperature = self.temperature() - if libnvml.nvmlCheckReturn(temperature, int): - temperature = str(temperature) + 'C' - return temperature + return f'{temperature}C' if libnvml.nvmlCheckReturn(temperature, int) else NA - def memory_loading_intensity(self): + def memory_loading_intensity(self) -> LoadingIntensity: return self.loading_intensity_of(self.memory_percent(), type='memory') - def gpu_loading_intensity(self): + def gpu_loading_intensity(self) -> LoadingIntensity: return self.loading_intensity_of(self.gpu_utilization(), type='gpu') - def loading_intensity(self): - loading_intensity = (self.memory_loading_intensity(), self.gpu_loading_intensity()) - if 'heavy' in loading_intensity: - return 'heavy' - if 'moderate' in loading_intensity: - return 'moderate' - return 'light' + def loading_intensity(self) -> LoadingIntensity: + return max(self.memory_loading_intensity(), self.gpu_loading_intensity()) - def display_color(self): + def display_color(self) -> str: if self.name().startswith('ERROR:'): return 'red' - return self.INTENSITY2COLOR.get(self.loading_intensity()) + return self.loading_intensity().color() - def memory_display_color(self): + def memory_display_color(self) -> str: if self.name().startswith('ERROR:'): return 'red' - return self.INTENSITY2COLOR.get(self.memory_loading_intensity()) + return self.memory_loading_intensity().color() - def gpu_display_color(self): + def gpu_display_color(self) -> str: if self.name().startswith('ERROR:'): return 'red' - return self.INTENSITY2COLOR.get(self.gpu_loading_intensity()) + return self.gpu_loading_intensity().color() @staticmethod - def loading_intensity_of(utilization, type='memory'): # pylint: disable=redefined-builtin + def loading_intensity_of( + utilization: float | str, + type: Literal['memory', 'gpu'] = 'memory', # pylint: disable=redefined-builtin + ) -> LoadingIntensity: thresholds = { 'memory': Device.MEMORY_UTILIZATION_THRESHOLDS, 'gpu': Device.GPU_UTILIZATION_THRESHOLDS, - }.get(type) + }[type] if utilization is NA: - return 'moderate' + return LoadingIntensity.MODERATE if isinstance(utilization, str): utilization = utilization.replace('%', '') utilization = float(utilization) if utilization >= thresholds[-1]: - return 'heavy' + return LoadingIntensity.HEAVY if utilization >= thresholds[0]: - return 'moderate' - return 'light' + return LoadingIntensity.MODERATE + return LoadingIntensity.LIGHT @staticmethod - def color_of(utilization, type='memory'): # pylint: disable=redefined-builtin - return Device.INTENSITY2COLOR.get(Device.loading_intensity_of(utilization, type=type)) + def color_of( + utilization: float | str, + type: Literal['memory', 'gpu'] = 'memory', # pylint: disable=redefined-builtin + ) -> str: + return Device.loading_intensity_of(utilization, type=type).color() -class MigDevice(MigDeviceBase, Device): - def __init__(self, *args, **kwargs): +class MigDevice(MigDeviceBase, Device): # type: ignore[misc] + def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - self._snapshot = None - self.tuple_index = (self.index,) if isinstance(self.index, int) else self.index - self.display_index = ':'.join(map(str, self.tuple_index)) + self._snapshot: Snapshot | None = None + self.tuple_index: tuple[int] | tuple[int, int] = ( + (self.index,) if isinstance(self.index, int) else self.index + ) + self.display_index: str = ':'.join(map(str, self.tuple_index)) def memory_usage(self) -> str: # string of used memory over total memory (in human readable) return f'{self.memory_used_human()} / {self.memory_total_human():>8s}' loading_intensity = Device.memory_loading_intensity - SNAPSHOT_KEYS = [ + SNAPSHOT_KEYS: ClassVar[list[str]] = [ 'name', 'memory_used', 'memory_free', diff --git a/nvitop/tui/library/displayable.py b/nvitop/tui/library/displayable.py index 8787aee..da0ee03 100644 --- a/nvitop/tui/library/displayable.py +++ b/nvitop/tui/library/displayable.py @@ -4,9 +4,22 @@ # pylint: disable=missing-module-docstring,missing-function-docstring +from __future__ import annotations + +from typing import TYPE_CHECKING, Generic, TypeVar + from nvitop.tui.library.libcurses import CursesShortcuts +if TYPE_CHECKING: + import curses + + from nvitop.tui.library.mouse import MouseEvent + + +__all__ = ['Displayable', 'DisplayableContainer'] + + class Displayable(CursesShortcuts): # pylint: disable=too-many-instance-attributes """Displayables are objects which are displayed on the screen. @@ -39,45 +52,47 @@ class Displayable(CursesShortcuts): # pylint: disable=too-many-instance-attribu x, y, width, height -- absolute coordinates and boundaries """ - def __init__(self, win, root=None): + def __init__(self, win: curses.window | None, root: DisplayableContainer | None = None) -> None: super().__init__() - self._need_redraw = True - self.focused = False - self._old_visible = self._visible = True - self.x = 0 - self.y = 0 - self._width = 0 - self.height = 0 + self._need_redraw: bool = True + self.focused: bool = False + self._old_visible: bool = True + self._visible: bool = True + self.x: int = 0 + self.y: int = 0 + self._width: int = 0 + self.height: int = 0 - self.win = win - self.root = root - self.parent = None + self.win: curses.window | None = win + self.root: DisplayableContainer | None = root + self.parent: DisplayableContainer | None = None - def __contains__(self, item): + def __contains__(self, item: Displayable | MouseEvent | tuple[int, int]) -> bool: """Check if item is inside the boundaries. item can be an iterable like [y, x] or an object with x and y methods. """ try: - y, x = item.y, item.x + y, x = item.y, item.x # type: ignore[union-attr] except AttributeError: try: - y, x = item + y, x = item # type: ignore[misc] except (ValueError, TypeError): return False return self.contains_point(y, x) - def contains_point(self, y, x): + def contains_point(self, y: int, x: int) -> bool: """Test whether the point lies inside this object. x and y should be absolute coordinates. """ return (self.x <= x < self.x + self.width) and (self.y <= y < self.y + self.height) - def poke(self): + def poke(self) -> None: """Called before drawing, even if invisible.""" + assert self.win is not None if self._old_visible != self.visible: self._old_visible = self.visible self.need_redraw = True @@ -85,44 +100,48 @@ class Displayable(CursesShortcuts): # pylint: disable=too-many-instance-attribu if not self.visible: self.win.erase() - def draw(self): + def draw(self) -> None: """Draw the object. Called on every main iteration if visible. Containers should call draw() on their contained objects here. Override this! """ + assert self.win is not None self.need_redraw = False - def finalize(self): + def finalize(self) -> None: """Called after every displayable is done drawing. Override this! """ + assert self.win is not None self.need_redraw = False - def destroy(self): + def destroy(self) -> None: """Called when the object is destroyed.""" self.win = None self.root = None - def click(self, event): + def click(self, event: MouseEvent) -> bool: # pylint: disable=unused-argument """Called when a mouse key is pressed and self.focused is True. Override this! """ + return False - def press(self, key): + def press(self, key: int) -> bool: # pylint: disable=unused-argument """Called when a key is pressed and self.focused is True. Override this! """ + return False @property - def visible(self): + def visible(self) -> bool: return self._visible @visible.setter - def visible(self, value): + def visible(self, value: bool) -> None: if self._visible != value: self.need_redraw = True self._visible = value @@ -130,31 +149,34 @@ class Displayable(CursesShortcuts): # pylint: disable=too-many-instance-attribu self.focused = False @property - def need_redraw(self): + def need_redraw(self) -> bool: return self._need_redraw @need_redraw.setter - def need_redraw(self, value): + def need_redraw(self, value: bool) -> None: if self._need_redraw != value: self._need_redraw = value if value and self.parent is not None and not self.parent.need_redraw: self.parent.need_redraw = True @property - def width(self): + def width(self) -> int: return self._width @width.setter - def width(self, value): + def width(self, value: int) -> None: if self.width != value and self.visible: self.need_redraw = True self._width = value - def __str__(self): + def __str__(self) -> str: return self.__class__.__name__ -class DisplayableContainer(Displayable): +D = TypeVar('D', bound=Displayable) + + +class DisplayableContainer(Displayable, Generic[D]): """DisplayableContainers are Displayables which contain other Displayables. This is also an abstract class. The methods draw, poke, finalize, @@ -172,20 +194,19 @@ class DisplayableContainer(Displayable): container -- a list with all contained objects (rw) """ - def __init__(self, win, root=None): + def __init__(self, win: curses.window | None, root: DisplayableContainer | None = None) -> None: super().__init__(win, root) - - self.container = [] + self.container: list[D] = [] # extended or overridden methods - def poke(self): + def poke(self) -> None: """Recursively called on objects in container.""" super().poke() for displayable in self.container: displayable.poke() - def draw(self): + def draw(self) -> None: """Recursively called on visible objects in container.""" for displayable in self.container: if self.need_redraw: @@ -195,19 +216,19 @@ class DisplayableContainer(Displayable): self.need_redraw = False - def finalize(self): + def finalize(self) -> None: """Recursively called on visible objects in container.""" for displayable in self.container: if displayable.visible: displayable.finalize() - def destroy(self): + def destroy(self) -> None: """Recursively called on objects in container.""" for displayable in self.container: displayable.destroy() super().destroy() - def press(self, key): + def press(self, key: int) -> bool: """Recursively called on objects in container.""" focused_obj = self.get_focused_obj() @@ -216,7 +237,7 @@ class DisplayableContainer(Displayable): return True return False - def click(self, event): + def click(self, event: MouseEvent) -> bool: """Recursively called on objects in container.""" focused_obj = self.get_focused_obj() if focused_obj and focused_obj.click(event): @@ -229,7 +250,7 @@ class DisplayableContainer(Displayable): # new methods - def add_child(self, obj): + def add_child(self, obj: D) -> None: """Add the objects to the container.""" if obj.parent is not None: obj.parent.remove_child(obj) @@ -237,13 +258,13 @@ class DisplayableContainer(Displayable): obj.parent = self obj.root = self.root - def replace_child(self, old_obj, new_obj): + def replace_child(self, old_obj: D, new_obj: D) -> None: """Replace the old object with the new instance in the container.""" self.container[self.container.index(old_obj)] = new_obj new_obj.parent = self new_obj.root = self.root - def remove_child(self, obj): + def remove_child(self, obj: D) -> None: """Remove the object from the container.""" try: self.container.remove(obj) @@ -253,13 +274,13 @@ class DisplayableContainer(Displayable): obj.parent = None obj.root = None - def get_focused_obj(self): + def get_focused_obj(self) -> D | None: # Finds a focused displayable object in the container. for displayable in self.container: if displayable.focused: return displayable try: - obj = displayable.get_focused_obj() + obj = displayable.get_focused_obj() # type: ignore[attr-defined] except AttributeError: pass else: diff --git a/nvitop/tui/library/history.py b/nvitop/tui/library/history.py index 583a61d..cecb7fa 100644 --- a/nvitop/tui/library/history.py +++ b/nvitop/tui/library/history.py @@ -3,26 +3,41 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import functools import itertools import threading import time from collections import deque +from typing import TYPE_CHECKING, Any, TypeVar from nvitop.api import NA -BOUND_UPDATE_INTERVAL = 1.0 +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + from typing_extensions import ParamSpec # Python 3.11+ + + _P = ParamSpec('_P') + _T = TypeVar('_T') + _S = TypeVar('_S') + + +__all__ = ['BufferedHistoryGraph', 'HistoryGraph'] + + +BOUND_UPDATE_INTERVAL: float = 1.0 # fmt: off -VALUE2SYMBOL_UP = { +VALUE2SYMBOL_UP: dict[tuple[int, int], str] = { (0, 0): ' ', (0, 1): '⢀', (0, 2): '⢠', (0, 3): '⢰', (0, 4): '⢸', (1, 0): '⡀', (1, 1): '⣀', (1, 2): '⣠', (1, 3): '⣰', (1, 4): '⣸', (2, 0): '⡄', (2, 1): '⣄', (2, 2): '⣤', (2, 3): '⣴', (2, 4): '⣼', (3, 0): '⡆', (3, 1): '⣆', (3, 2): '⣦', (3, 3): '⣶', (3, 4): '⣾', (4, 0): '⡇', (4, 1): '⣇', (4, 2): '⣧', (4, 3): '⣷', (4, 4): '⣿', } -VALUE2SYMBOL_DOWN = { +VALUE2SYMBOL_DOWN: dict[tuple[int, int], str] = { (0, 0): ' ', (0, 1): '⠈', (0, 2): '⠘', (0, 3): '⠸', (0, 4): '⢸', (1, 0): '⠁', (1, 1): '⠉', (1, 2): '⠙', (1, 3): '⠹', (1, 4): '⢹', (2, 0): '⠃', (2, 1): '⠋', (2, 2): '⠛', (2, 3): '⠻', (2, 4): '⢻', @@ -30,23 +45,28 @@ VALUE2SYMBOL_DOWN = { (4, 0): '⡇', (4, 1): '⡏', (4, 2): '⡟', (4, 3): '⡿', (4, 4): '⣿', } # fmt: on -SYMBOL2VALUE_UP = {v: k for k, v in VALUE2SYMBOL_UP.items()} -SYMBOL2VALUE_DOWN = {v: k for k, v in VALUE2SYMBOL_DOWN.items()} -PAIR2SYMBOL_UP = { +SYMBOL2VALUE_UP: dict[str, tuple[int, int]] = {v: k for k, v in VALUE2SYMBOL_UP.items()} +SYMBOL2VALUE_DOWN: dict[str, tuple[int, int]] = {v: k for k, v in VALUE2SYMBOL_DOWN.items()} +PAIR2SYMBOL_UP: dict[tuple[str, str], str] = { (s1, s2): VALUE2SYMBOL_UP[SYMBOL2VALUE_UP[s1][-1], SYMBOL2VALUE_UP[s2][0]] for s1, s2 in itertools.product(SYMBOL2VALUE_UP, repeat=2) } -PAIR2SYMBOL_DOWN = { +PAIR2SYMBOL_DOWN: dict[tuple[str, str], str] = { (s1, s2): VALUE2SYMBOL_DOWN[SYMBOL2VALUE_DOWN[s1][-1], SYMBOL2VALUE_DOWN[s2][0]] for s1, s2 in itertools.product(SYMBOL2VALUE_DOWN, repeat=2) } -GRAPH_SYMBOLS = ''.join( +GRAPH_SYMBOLS: str = ''.join( sorted(set(itertools.chain(VALUE2SYMBOL_UP.values(), VALUE2SYMBOL_DOWN.values()))), ).replace(' ', '') -def grouped(iterable, size, fillvalue=None): - yield from itertools.zip_longest(*([iter(iterable)] * size), fillvalue=fillvalue) +def grouped( + iterable: Iterable[_T], + size: int, + fillvalue: _S = None, # type: ignore[assignment] +) -> itertools.zip_longest[tuple[_T | _S, ...]]: + it = iter(iterable) + return itertools.zip_longest(*([it] * size), fillvalue=fillvalue) class HistoryGraph: # pylint: disable=too-many-instance-attributes @@ -55,24 +75,24 @@ class HistoryGraph: # pylint: disable=too-many-instance-attributes # pylint: disable-next=too-many-arguments def __init__( self, - upperbound, - width, - height, + upperbound: float, + width: int, + height: int, *, - format='{:.1f}'.format, # pylint: disable=redefined-builtin - max_format=None, - baseline=0.0, - dynamic_bound=False, - min_bound=None, - init_bound=None, - upsidedown=False, - ): + format: Callable[[float], str] = '{:.1f}'.format, # pylint: disable=redefined-builtin + max_format: Callable[[float], str] | None = None, + baseline: float = 0.0, + dynamic_bound: bool = False, + min_bound: float | None = None, + init_bound: float | None = None, + upsidedown: bool = False, + ) -> None: assert baseline < upperbound - self.format = format + self.format: Callable[[float], str] = format if max_format is None: max_format = format - self.max_format = max_format + self.max_format: Callable[[float], str] = max_format if dynamic_bound: if min_bound is None: @@ -83,29 +103,35 @@ class HistoryGraph: # pylint: disable=too-many-instance-attributes assert min_bound is None assert init_bound is None min_bound = init_bound = upperbound - self.baseline = baseline - self.min_bound = min_bound - self.max_bound = upperbound - self.bound = init_bound - self.next_bound_update_at = time.monotonic() - self._width = width - self._height = height + self.baseline: float = baseline + self.min_bound: float = min_bound + self.max_bound: float = upperbound + self.bound: float = init_bound + self.next_bound_update_at: float = time.monotonic() + self._width: int = width + self._height: int = height - self.maxlen = 2 * self.width + 1 - self.history = deque( + self.maxlen: int = 2 * self.width + 1 + self.history: deque[float] = deque( [self.baseline - 0.1] * (2 * self.MAX_WIDTH + 1), maxlen=(2 * self.MAX_WIDTH + 1), ) - self.reversed_history = deque([self.baseline - 0.1] * self.maxlen, maxlen=self.maxlen) - self._max_value_maintainer = deque([self.baseline - 0.1] * self.maxlen, maxlen=self.maxlen) - self.last_retval = None + self.reversed_history: deque[float] = deque( + [self.baseline - 0.1] * self.maxlen, + maxlen=self.maxlen, + ) + self._max_value_maintainer: deque[float] = deque( + [self.baseline - 0.1] * self.maxlen, + maxlen=self.maxlen, + ) + self.last_retval: Any = None - self.graph = [] - self.last_graph = [] - self.upsidedown = upsidedown + self.graph: list[str] = [] + self.last_graph: list[str] = [] + self.upsidedown: bool = upsidedown if upsidedown: - self.value2symbol = VALUE2SYMBOL_DOWN - self.pair2symbol = PAIR2SYMBOL_DOWN + self.value2symbol: dict[tuple[int, int], str] = VALUE2SYMBOL_DOWN + self.pair2symbol: dict[tuple[str, str], str] = PAIR2SYMBOL_DOWN else: self.value2symbol = VALUE2SYMBOL_UP self.pair2symbol = PAIR2SYMBOL_UP @@ -115,11 +141,11 @@ class HistoryGraph: # pylint: disable=too-many-instance-attributes self.remake_graph() @property - def width(self): + def width(self) -> int: return self._width @width.setter - def width(self, value): + def width(self, value: int) -> None: if self._width != value: assert isinstance(value, int) assert value >= 1 @@ -136,7 +162,7 @@ class HistoryGraph: # pylint: disable=too-many-instance-attributes ) for history in itertools.islice( self.history, - max(0, self.history.maxlen - self.maxlen), + max(0, self.history.maxlen - self.maxlen), # type: ignore[operator] self.history.maxlen, ): if self.reversed_history[-1] == self._max_value_maintainer[0]: @@ -151,11 +177,11 @@ class HistoryGraph: # pylint: disable=too-many-instance-attributes self.remake_graph() @property - def height(self): + def height(self) -> int: return self._height @height.setter - def height(self, value): + def height(self, value: int) -> None: if self._height != value: assert isinstance(value, int) assert value >= 1 @@ -163,11 +189,11 @@ class HistoryGraph: # pylint: disable=too-many-instance-attributes self.remake_graph() @property - def graph_size(self): + def graph_size(self) -> tuple[int, int]: return (self.width, self.height) @graph_size.setter - def graph_size(self, value): + def graph_size(self, value: tuple[int, int]) -> None: width, height = value assert isinstance(width, int) assert width >= 1 @@ -178,38 +204,37 @@ class HistoryGraph: # pylint: disable=too-many-instance-attributes self.width = width @property - def last_value(self): + def last_value(self) -> float: return self.reversed_history[0] @property - def max_value(self): + def max_value(self) -> float: return self._max_value_maintainer[0] - def last_value_string(self): + def last_value_string(self) -> str: last_value = self.last_value if last_value >= self.baseline: return self.format(last_value) try: - return self.format(NA) + return self.format(NA) # type: ignore[arg-type] except ValueError: return NA __str__ = last_value_string - def max_value_string(self): + def max_value_string(self) -> str: max_value = self.max_value if max_value >= self.baseline: return self.max_format(max_value) try: - return self.max_format(NA) + return self.max_format(NA) # type: ignore[arg-type] except ValueError: return NA - def add(self, value): - if value is NA: + def add(self, value: float) -> None: + if value is NA: # type: ignore[comparison-overlap] value = self.baseline - 0.1 - if not isinstance(value, (int, float)): - return + assert isinstance(value, (int, float)) with self.write_lock: if self.reversed_history[-1] == self._max_value_maintainer[0]: @@ -234,7 +259,7 @@ class HistoryGraph: # pylint: disable=too-many-instance-attributes for i, (line, char) in enumerate(zip(self.graph, bar)): self.graph[i] = (line + char)[-self.width :] - def remake_graph(self): + def remake_graph(self) -> None: with self.remake_lock: if self.max_value >= self.baseline: reversed_bars = [] @@ -254,7 +279,7 @@ class HistoryGraph: # pylint: disable=too-many-instance-attributes self.graph = [' ' * self.width for _ in range(self.height)] self.last_graph = [' ' * (self.width - 1) for _ in range(self.height)] - def make_bar(self, value1, value2): + def make_bar(self, value1: float, value2: float) -> list[str]: if self.bound <= self.baseline: return [' '] * self.height @@ -274,22 +299,26 @@ class HistoryGraph: # pylint: disable=too-many-instance-attributes bar.reverse() return bar - def shift_line(self, line): - return ''.join(map(self.pair2symbol.get, zip(line, line[1:]))) + def shift_line(self, line: str) -> str: + return ''.join(self.pair2symbol[p] for p in zip(line, line[1:])) - def __getitem__(self, item): + def __getitem__(self, item: int) -> float: return self.reversed_history[item] - def hook(self, func, get_value=None): + def hook( + self, + func: Callable[_P, _T], + *, + get_value: Callable[[_T], float] | None = None, + ) -> Callable[_P, _T]: @functools.wraps(func) - def wrapped(*args, **kwargs): - self.last_retval = retval = value = func(*args, **kwargs) - if get_value is not None: - value = get_value(retval) + def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _T: + self.last_retval = retval = func(*args, **kwargs) + value: float = get_value(retval) if get_value is not None else retval # type: ignore[assignment] self.add(value) return retval - wrapped.history = self + wrapped.history = self # type: ignore[attr-defined] return wrapped __call__ = hook @@ -299,19 +328,19 @@ class BufferedHistoryGraph(HistoryGraph): # pylint: disable-next=too-many-arguments def __init__( self, - upperbound, - width, - height, + upperbound: float, + width: int, + height: int, *, - format='{:.1f}'.format, # pylint: disable=redefined-builtin - max_format=None, - baseline=0.0, - dynamic_bound=False, - upsidedown=False, - min_bound=None, - init_bound=None, - interval=1.0, - ): + format: Callable[[float], str] = '{:.1f}'.format, # pylint: disable=redefined-builtin + max_format: Callable[[float], str] | None = None, + baseline: float = 0.0, + dynamic_bound: bool = False, + min_bound: float | None = None, + init_bound: float | None = None, + upsidedown: bool = False, + interval: float = 1.0, + ) -> None: assert interval > 0.0 super().__init__( upperbound, @@ -326,23 +355,22 @@ class BufferedHistoryGraph(HistoryGraph): upsidedown=upsidedown, ) - self.interval = interval - self.start_time = time.monotonic() - self.last_update_time = self.start_time - self.buffer = [] + self.interval: float = interval + self.start_time: float = time.monotonic() + self.last_update_time: float = self.start_time + self.buffer: list[float] = [] @property - def last_value(self): + def last_value(self) -> float: last_value = super().last_value if last_value < self.baseline and len(self.buffer) > 0: return sum(self.buffer) / len(self.buffer) return last_value - def add(self, value): - if value is NA: + def add(self, value: float) -> None: + if value is NA: # type: ignore[comparison-overlap] value = self.baseline - 0.1 - if not isinstance(value, (int, float)): - return + assert isinstance(value, (int, float)) timestamp = time.monotonic() timedelta = timestamp - self.last_update_time diff --git a/nvitop/tui/library/host.py b/nvitop/tui/library/host.py index b92e892..ada9888 100644 --- a/nvitop/tui/library/host.py +++ b/nvitop/tui/library/host.py @@ -3,16 +3,16 @@ # pylint: disable=missing-module-docstring,missing-function-docstring +from __future__ import annotations + from types import MappingProxyType -from typing import TYPE_CHECKING, NamedTuple +from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar from nvitop.api import NA, host -from nvitop.api.host import WINDOWS, WSL, AccessDenied, PsutilError +from nvitop.api.host import AccessDenied, PsutilError __all__ = [ - 'WINDOWS', - 'WSL', 'AccessDenied', 'PsutilError', 'cpu_percent', @@ -26,11 +26,19 @@ __all__ = [ ] -def ignore_error(*, fallback): +if TYPE_CHECKING: + from collections.abc import Callable + from typing_extensions import ParamSpec # Python 3.10+ + + _P = ParamSpec('_P') + _T = TypeVar('_T') + + +def ignore_error(*, fallback: Any) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: """Ignore errors in the function.""" - def wrapper(func): - def wrapped(*args, **kwargs): + def wrapper(func: Callable[_P, _T]) -> Callable[_P, _T]: + def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _T: try: return func(*args, **kwargs) except Exception: # noqa: BLE001 # pylint: disable=broad-exception-caught @@ -42,15 +50,15 @@ def ignore_error(*, fallback): class VirtualMemory(NamedTuple): # pylint: disable=missing-class-docstring - total: int - available: int - percent: int - used: int - free: int + total: int = NA # type: ignore[assignment] + available: int = NA # type: ignore[assignment] + percent: int = NA # type: ignore[assignment] + used: int = NA # type: ignore[assignment] + free: int = NA # type: ignore[assignment] -@ignore_error(fallback=VirtualMemory(NA, NA, NA, NA, NA)) -def virtual_memory(): +@ignore_error(fallback=VirtualMemory()) +def virtual_memory() -> VirtualMemory: vm = host.virtual_memory() return VirtualMemory( total=vm.total, @@ -62,32 +70,28 @@ def virtual_memory(): class SwapMemory(NamedTuple): # pylint: disable=missing-class-docstring - total: int - used: int - free: int - percent: float - sin: int - sout: int + total: int = NA # type: ignore[assignment] + used: int = NA # type: ignore[assignment] + free: int = NA # type: ignore[assignment] + percent: float = NA # type: ignore[assignment] -@ignore_error(fallback=SwapMemory(NA, NA, NA, NA, NA, NA)) -def swap_memory(): +@ignore_error(fallback=SwapMemory()) +def swap_memory() -> SwapMemory: sm = host.swap_memory() return SwapMemory( total=sm.total, used=sm.used, free=sm.free, percent=sm.percent, - sin=sm.sin, - sout=sm.sout, ) @ignore_error(fallback=(NA, NA, NA)) -def load_average(): +def load_average() -> tuple[float, float, float]: la = host.load_average() if la is None: - return (NA, NA, NA) + return (NA, NA, NA) # type: ignore[unreachable] return la diff --git a/nvitop/tui/library/keybinding.py b/nvitop/tui/library/keybinding.py index 9d8d392..50f9aff 100644 --- a/nvitop/tui/library/keybinding.py +++ b/nvitop/tui/library/keybinding.py @@ -4,19 +4,52 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import copy import curses import curses.ascii import string from collections import OrderedDict +from typing import TYPE_CHECKING, Callable, Dict, Tuple, Union -DIGITS = set(map(ord, string.digits)) +if TYPE_CHECKING: + from collections.abc import Generator + from typing_extensions import TypeAlias # Python 3.10+ + + +__all__ = [ + 'ALT_KEY', + 'ANYKEY', + 'DIGITS', + 'PASSIVE_ACTION', + 'QUANT_KEY', + 'REVERSED_SPECIAL_KEYS', + 'SPECIAL_KEYS', + 'SPECIAL_KEYS', + 'SPECIAL_KEYS_UNCASED', + 'KeyBuffer', + 'KeyMaps', + 'construct_keybinding', + 'normalize_keybinding', + 'parse_keybinding', +] + + +IntKey: TypeAlias = Union[int, Tuple[int, ...]] + + +DIGITS: frozenset[int] = frozenset(map(ord, string.digits)) # Arbitrary numbers which are not used with curses.KEY_XYZ -ANYKEY, PASSIVE_ACTION, ALT_KEY, QUANT_KEY = range(9001, 9005) +ANYKEY: int = 9001 +PASSIVE_ACTION: int = 9002 +ALT_KEY: int = 9003 +QUANT_KEY: int = 9004 -SPECIAL_KEYS = OrderedDict( + +NAMED_SPECIAL_KEYS: OrderedDict[str, int] = OrderedDict( [ ('BS', curses.KEY_BACKSPACE), ('Backspace', curses.KEY_BACKSPACE), # overrides in REVERSED_SPECIAL_KEYS @@ -44,18 +77,12 @@ SPECIAL_KEYS = OrderedDict( ('gt', ord('>')), ], ) - -NAMED_SPECIAL_KEYS = tuple(SPECIAL_KEYS.keys()) -SPECIAL_KEYS_UNCASED = {} -VERY_SPECIAL_KEYS = { - 'Alt': ALT_KEY, - 'any': ANYKEY, - 'bg': PASSIVE_ACTION, - 'allow_quantifiers': QUANT_KEY, -} +SPECIAL_KEYS: OrderedDict[str, IntKey] +SPECIAL_KEYS_UNCASED: dict[str, IntKey] +REVERSED_SPECIAL_KEYS: dict[IntKey, str] -def _uncase_special_key(key_string): +def _uncase_special_key(key_string: str) -> str: """Uncase a special key. >>> _uncase_special_key('Esc') @@ -77,65 +104,78 @@ def _uncase_special_key(key_string): return uncased -def _special_keys_init(): - for key, val in tuple(SPECIAL_KEYS.items()): - SPECIAL_KEYS['M-' + key] = (ALT_KEY, val) - SPECIAL_KEYS['A-' + key] = (ALT_KEY, val) # overrides in REVERSED_SPECIAL_KEYS +def _special_keys_init() -> None: + global SPECIAL_KEYS, SPECIAL_KEYS_UNCASED, REVERSED_SPECIAL_KEYS # pylint: disable=global-statement + + SPECIAL_KEYS = NAMED_SPECIAL_KEYS.copy() # type: ignore[assignment] + for key, int_value in NAMED_SPECIAL_KEYS.items(): + SPECIAL_KEYS[f'M-{key}'] = (ALT_KEY, int_value) + SPECIAL_KEYS[f'A-{key}'] = (ALT_KEY, int_value) # overrides in REVERSED_SPECIAL_KEYS for char in 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_!{}[],./': - SPECIAL_KEYS['M-' + char] = (ALT_KEY, ord(char)) - SPECIAL_KEYS['A-' + char] = (ALT_KEY, ord(char)) # overrides in REVERSED_SPECIAL_KEYS + SPECIAL_KEYS[f'M-{char}'] = (ALT_KEY, ord(char)) + SPECIAL_KEYS[f'A-{char}'] = (ALT_KEY, ord(char)) # overrides in REVERSED_SPECIAL_KEYS # We will need to reorder the keys of SPECIAL_KEYS below. # For example, will override in REVERSE_SPECIAL_KEYS, # this makes construct_keybinding(parse_keybinding('')) == '' for char in 'abcdefghijklmnopqrstuvwxyz_': - SPECIAL_KEYS['C-' + char] = ord(char) - 96 - + SPECIAL_KEYS[f'C-{char}'] = ord(char) - 96 SPECIAL_KEYS['C-Space'] = 0 for n in range(64): - SPECIAL_KEYS['F' + str(n)] = curses.KEY_F0 + n + SPECIAL_KEYS[f'F{n}'] = curses.KEY_F0 + n - SPECIAL_KEYS.update(VERY_SPECIAL_KEYS) # noqa: F821 + # Very special keys + SPECIAL_KEYS.update( + { + 'Alt': ALT_KEY, + 'any': ANYKEY, + 'bg': PASSIVE_ACTION, + 'allow_quantifiers': QUANT_KEY, + }, + ) - # Reorder the keys of SPECIAL_KEYS. - for key in NAMED_SPECIAL_KEYS: # noqa: F821 + # Reorder the keys of SPECIAL_KEYS + for key in NAMED_SPECIAL_KEYS: SPECIAL_KEYS.move_to_end(key, last=True) - for key, val in SPECIAL_KEYS.items(): - SPECIAL_KEYS_UNCASED[_uncase_special_key(key)] = val + SPECIAL_KEYS_UNCASED = {_uncase_special_key(k): v for k, v in SPECIAL_KEYS.items()} + REVERSED_SPECIAL_KEYS = {v: k for k, v in SPECIAL_KEYS.items()} _special_keys_init() -del _special_keys_init, VERY_SPECIAL_KEYS, NAMED_SPECIAL_KEYS -REVERSED_SPECIAL_KEYS = OrderedDict([(v, k) for k, v in SPECIAL_KEYS.items()]) +del _special_keys_init -def parse_keybinding(obj): # pylint: disable=too-many-branches +def parse_keybinding(obj: IntKey | str) -> tuple[int, ...]: # pylint: disable=too-many-branches r"""Translate a keybinding to a sequence of integers The letter case of special keys in the keybinding string will be ignored. - >>> out = tuple(parse_keybinding('lol')) + >>> out = parse_keybinding('lol') >>> out (108, 111, 108, 10) >>> out == (ord('l'), ord('o'), ord('l'), ord('\n')) True - >>> out = tuple(parse_keybinding('x')) + >>> out = parse_keybinding('x') >>> out (120, 9003, 260) >>> out == (ord('x'), ALT_KEY, curses.KEY_LEFT) True """ assert isinstance(obj, (tuple, int, str)) - if isinstance(obj, tuple): - yield from obj - elif isinstance(obj, int): # pylint: disable=too-many-nested-blocks - yield obj - else: # pylint: disable=too-many-nested-blocks + + def parse(obj: IntKey | str) -> Generator[int]: # pylint: disable=too-many-branches + if isinstance(obj, tuple): + yield from obj + return + if isinstance(obj, int): + yield obj + return + in_brackets = False - bracket_content = [] + bracket_content: list[str] = [] for char in obj: if in_brackets: if char == '>': @@ -143,7 +183,7 @@ def parse_keybinding(obj): # pylint: disable=too-many-branches key_string = ''.join(bracket_content) try: keys = SPECIAL_KEYS_UNCASED[_uncase_special_key(key_string)] - yield from keys + yield from keys # type: ignore[misc] except KeyError: if key_string.isdigit(): yield int(key_string) @@ -153,7 +193,7 @@ def parse_keybinding(obj): # pylint: disable=too-many-branches yield ord(bracket_char) yield ord('>') except TypeError: - yield keys # it was no tuple, just an int + yield keys # type: ignore[misc] # it was not tuple, just an int else: bracket_content.append(char) elif char == '<': @@ -166,16 +206,16 @@ def parse_keybinding(obj): # pylint: disable=too-many-branches for char in bracket_content: yield ord(char) + return tuple(parse(obj)) -def key_to_string(key): + +def key_to_string(key: IntKey) -> str: if key in range(33, 127): - return chr(key) - if key in REVERSED_SPECIAL_KEYS: - return f'<{REVERSED_SPECIAL_KEYS[key]}>' - return f'<{key}>' + return chr(key) # type: ignore[arg-type] + return f'<{REVERSED_SPECIAL_KEYS.get(key, key)}>' -def construct_keybinding(keys): +def construct_keybinding(keys: IntKey) -> str: """Do the reverse of parse_keybinding. >>> construct_keybinding(parse_keybinding('lol')) @@ -187,11 +227,7 @@ def construct_keybinding(keys): >>> construct_keybinding(parse_keybinding('x')) 'x' """ - try: - keys = tuple(keys) - except TypeError: - assert isinstance(keys, int) - keys = (keys,) + keys = (keys,) if isinstance(keys, int) else tuple(keys) strings = [] alt_key_on = False for key in keys: @@ -210,7 +246,7 @@ def construct_keybinding(keys): return ''.join(strings) -def normalize_keybinding(keybinding): +def normalize_keybinding(keybinding: str) -> str: """Normalize a keybinding to a string. >>> normalize_keybinding('lol') @@ -225,113 +261,125 @@ def normalize_keybinding(keybinding): return construct_keybinding(parse_keybinding(keybinding)) -class KeyMaps(dict): - def __init__(self, keybuffer=None): - super().__init__() - self.keybuffer = keybuffer - self.used_keymap = None +KeyMapPointer: TypeAlias = Dict[int, Union['KeyMapPointer', Callable[[], None]]] - def use_keymap(self, keymap_name): + +class KeyMaps(Dict[str, KeyMapPointer]): + def __init__(self, keybuffer: KeyBuffer) -> None: + super().__init__() + self.keybuffer: KeyBuffer = keybuffer + self.used_keymap: str | None = None + + def use_keymap(self, keymap_name: str) -> None: self.keybuffer.keymap = self.get(keymap_name, {}) if self.used_keymap != keymap_name: self.used_keymap = keymap_name self.keybuffer.clear() - def clear_keymap(self, keymap_name): - self[keymap_name] = {} + def clear_keymap(self, keymap_name: str) -> KeyMapPointer: + keymap = self[keymap_name] = {} if self.used_keymap == keymap_name: - self.keybuffer.keymap = {} + self.keybuffer.keymap = keymap self.keybuffer.clear() + return keymap - def _clean_input(self, context, keys): + def _clean_input(self, context: str, keybinding: str) -> tuple[tuple[int, ...], KeyMapPointer]: try: pointer = self[context] except KeyError: self[context] = pointer = {} - keys = keys.encode('utf-8').decode('latin-1') - return list(parse_keybinding(keys)), pointer + keybinding = keybinding.encode('utf-8').decode('latin-1') + return parse_keybinding(keybinding), pointer - def bind(self, context, keys, leaf): - keys, pointer = self._clean_input(context, keys) + def bind(self, context: str, keybinding: str, leaf: Callable[[], None]) -> None: + keys, pointer = self._clean_input(context, keybinding) if not keys: return last_key = keys[-1] for key in keys[:-1]: if key in pointer and isinstance(pointer[key], dict): - pointer = pointer[key] + pointer = pointer[key] # type: ignore[assignment] else: pointer = pointer[key] = {} pointer[last_key] = leaf - def copy(self, context, source, target): + def alias(self, context: str, source: str, target: str) -> None: clean_source, pointer = self._clean_input(context, source) if not source: return for key in clean_source: try: - pointer = pointer[key] + pointer = pointer[key] # type: ignore[assignment] except KeyError as ex: # noqa: PERF203 raise KeyError( f'Tried to copy the keybinding `{source}`, but it was not found.', ) from ex try: - self.bind(context, target, copy.deepcopy(pointer)) + self.bind(context, target, copy.deepcopy(pointer)) # type: ignore[arg-type] except TypeError: - self.bind(context, target, pointer) + self.bind(context, target, pointer) # type: ignore[arg-type] - def unbind(self, context, keys): - keys, pointer = self._clean_input(context, keys) + def unbind(self, context: str, keybinding: str) -> None: + keys, pointer = self._clean_input(context, keybinding) if not keys: return - self._unbind_traverse(pointer, keys) - @staticmethod - def _unbind_traverse(pointer, keys, pos=0): - if keys[pos] not in pointer: - return - if len(keys) > pos + 1 and isinstance(pointer, dict): - KeyMaps._unbind_traverse(pointer[keys[pos]], keys, pos=pos + 1) - if not pointer[keys[pos]]: - del pointer[keys[pos]] - elif len(keys) == pos + 1: - try: - del pointer[keys[pos]] - except KeyError: - pass - try: - keys.pop() - except IndexError: - pass + def unbind_traverse(pointer: KeyMapPointer, keys: list[int], pos: int = 0) -> None: + if keys[pos] not in pointer: + return + if len(keys) > pos + 1 and isinstance(pointer, dict): + unbind_traverse(pointer[keys[pos]], keys, pos=pos + 1) # type: ignore[arg-type] + if not pointer[keys[pos]]: + del pointer[keys[pos]] + elif len(keys) == pos + 1: + try: + del pointer[keys[pos]] + except KeyError: + pass + try: + keys.pop() + except IndexError: + pass + + unbind_traverse(pointer, list(keys)) class KeyBuffer: # pylint: disable=too-many-instance-attributes - any_key = ANYKEY - passive_key = PASSIVE_ACTION - quantifier_key = QUANT_KEY - excluded_from_anykey = [curses.ascii.ESC] + class QuantifierFinished: # pylint: disable=too-few-public-methods + pass - def __init__(self, keymap=None): - self.keymap = keymap - self.keys = [] - self.wildcards = [] - self.pointer = self.keymap - self.result = None - self.quantifier = None - self.finished_parsing_quantifier = False - self.finished_parsing = False - self.parse_error = False + QUANTIFIER_KEY_FINISHED = QuantifierFinished() + + del QuantifierFinished + + any_key: int = ANYKEY + passive_key: int = PASSIVE_ACTION + quantifier_key: int = QUANT_KEY + excluded_from_anykey: frozenset[int] = frozenset({curses.ascii.ESC}) + + def __init__(self, keymap: KeyMapPointer | None = None) -> None: + self.keymap: KeyMapPointer | None = keymap + self.keys: list[int] = [] + self.wildcards: list[int] = [] + self.pointer: KeyMapPointer | None = self.keymap + self.result: Callable[[], None] | None = None + self.quantifier: int | None = None + self.finished_parsing_quantifier: bool = False + self.finished_parsing: bool = False + self.parse_error: bool = False if ( self.keymap and self.quantifier_key in self.keymap - and self.keymap[self.quantifier_key] == 'false' + and self.keymap[self.quantifier_key] is self.QUANTIFIER_KEY_FINISHED # type: ignore[comparison-overlap] ): self.finished_parsing_quantifier = True - def clear(self): - self.__init__(self.keymap) # pylint: disable=unnecessary-dunder-call + def clear(self) -> None: + self.__init__(self.keymap) # type: ignore[misc] # pylint: disable=unnecessary-dunder-call - def add(self, key): + def add(self, key: int) -> None: + assert self.pointer is not None self.keys.append(key) self.result = None if not self.finished_parsing_quantifier and key in DIGITS: @@ -343,23 +391,29 @@ class KeyBuffer: # pylint: disable=too-many-instance-attributes moved = True if key in self.pointer: - self.pointer = self.pointer[key] + self.pointer = self.pointer[key] # type: ignore[assignment] elif self.any_key in self.pointer and key not in self.excluded_from_anykey: self.wildcards.append(key) - self.pointer = self.pointer[self.any_key] + self.pointer = self.pointer[self.any_key] # type: ignore[assignment] else: moved = False if moved: if isinstance(self.pointer, dict): if self.passive_key in self.pointer: - self.result = self.pointer[self.passive_key] + self.result = self.pointer[self.passive_key] # type: ignore[assignment] else: - self.result = self.pointer + self.result = self.pointer # type: ignore[unreachable] self.finished_parsing = True else: self.finished_parsing = True self.parse_error = True - def __str__(self): - return construct_keybinding(self.keys) + def __str__(self) -> str: + return construct_keybinding(tuple(self.keys)) + + +if __name__ == '__main__': + import doctest + + doctest.testmod() diff --git a/nvitop/tui/library/libcurses.py b/nvitop/tui/library/libcurses.py index f3fe963..cab2b78 100644 --- a/nvitop/tui/library/libcurses.py +++ b/nvitop/tui/library/libcurses.py @@ -3,47 +3,56 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import colorsys import contextlib import curses import locale import os import signal +from typing import TYPE_CHECKING, Any, ClassVar, Tuple, Union from nvitop.tui.library.history import GRAPH_SYMBOLS -LIGHT_THEME = False -DEFAULT_FOREGROUND = curses.COLOR_WHITE -DEFAULT_BACKGROUND = curses.COLOR_BLACK -COLOR_PAIRS = {None: 0} -TRUE_COLORS = dict( - [ - ('black', 0), - ('red', 1), - ('green', 2), - ('yellow', 3), - ('blue', 4), - ('magenta', 5), - ('cyan', 6), - ('white', 7), - ('bright black', 8), - ('bright red', 9), - ('bright green', 10), - ('bright yellow', 11), - ('bright blue', 12), - ('bright magenta', 13), - ('bright cyan', 14), - ('bright white', 15), - ] - + [(f'preserved {i:02d}', i) for i in range(16, 64)], -) +if TYPE_CHECKING: + from collections.abc import Generator + from typing_extensions import TypeAlias # Python 3.10+ -BASE_ATTR = 0 +__all__ = ['CursesShortcuts', 'libcurses', 'setlocale_utf8'] -def _init_color_theme(light_theme=False): +LIGHT_THEME: bool = False +DEFAULT_FOREGROUND: int = curses.COLOR_WHITE +DEFAULT_BACKGROUND: int = curses.COLOR_BLACK +COLOR_PAIRS: dict[tuple[int, int], int] = {} +TRUE_COLORS: dict[str | tuple[int, int, int], int] = { + 'black': 0, + 'red': 1, + 'green': 2, + 'yellow': 3, + 'blue': 4, + 'magenta': 5, + 'cyan': 6, + 'white': 7, + 'bright black': 8, + 'bright red': 9, + 'bright green': 10, + 'bright yellow': 11, + 'bright blue': 12, + 'bright magenta': 13, + 'bright cyan': 14, + 'bright white': 15, + **{f'preserved {i:02d}': i for i in range(16, 64)}, +} + + +BASE_ATTR: int = 0 + + +def _init_color_theme(light_theme: bool = False) -> None: """Set the default fg/bg colors.""" global LIGHT_THEME, DEFAULT_FOREGROUND, DEFAULT_BACKGROUND # pylint: disable=global-statement @@ -56,7 +65,7 @@ def _init_color_theme(light_theme=False): DEFAULT_BACKGROUND = curses.COLOR_BLACK -def _colormap(x, levels=160): +def _colormap(x: float, levels: int = 160) -> tuple[int, int, int]: # pylint: disable=invalid-name h = 0.5 * (1.0 - x) - 0.15 h = (round(h * levels) / levels) % 1.0 @@ -64,7 +73,7 @@ def _colormap(x, levels=160): return (round(1000.0 * r), round(1000.0 * g), round(1000.0 * b)) -def _get_true_color(rgb): +def _get_true_color(rgb: tuple[int, int, int]) -> int: if rgb not in TRUE_COLORS: try: curses.init_color(len(TRUE_COLORS), *rgb) @@ -74,7 +83,10 @@ def _get_true_color(rgb): return TRUE_COLORS[rgb] -def _get_color(fg, bg): +Color: TypeAlias = Union[str, int, float, Tuple[int, int, int]] + + +def _get_color(fg: Color, bg: Color) -> int: """Return the curses color pair for the given fg/bg combination.""" global COLOR_PAIRS # pylint: disable=global-statement,global-variable-not-assigned @@ -93,9 +105,9 @@ def _get_color(fg, bg): key = (fg, bg) if key not in COLOR_PAIRS: - size = len(COLOR_PAIRS) + new_id = len(COLOR_PAIRS) + 1 try: - curses.init_pair(size, fg, bg) + curses.init_pair(new_id, fg, bg) except curses.error: # If curses.use_default_colors() failed during the initialization # of curses, then using -1 as fg or bg will fail as well, which @@ -106,16 +118,16 @@ def _get_color(fg, bg): bg = DEFAULT_BACKGROUND try: - curses.init_pair(size, fg, bg) + curses.init_pair(new_id, fg, bg) except curses.error: # If this fails too, colors are probably not supported pass - COLOR_PAIRS[key] = size + COLOR_PAIRS[key] = new_id return COLOR_PAIRS[key] -def setlocale_utf8(): +def setlocale_utf8() -> bool: for code in ('C.UTF-8', 'en_US.UTF-8', '', 'C'): try: code = locale.setlocale(locale.LC_ALL, code) @@ -129,7 +141,7 @@ def setlocale_utf8(): @contextlib.contextmanager -def libcurses(colorful=False, light_theme=False): +def libcurses(colorful: bool = False, light_theme: bool = False) -> Generator[curses.window]: os.environ.setdefault('ESCDELAY', '25') setlocale_utf8() @@ -160,7 +172,7 @@ def libcurses(colorful=False, light_theme=False): pass # Push a Ctrl+C (ascii value 3) to the curses getch stack - def interrupt_handler(signalnum, frame): # pylint: disable=unused-argument + def interrupt_handler(*_: Any) -> None: # pylint: disable=unused-argument curses.ungetch(3) # Simulate a ^C press in curses when an interrupt is caught @@ -182,69 +194,80 @@ class CursesShortcuts: addstr(*args) -- failsafe version of self.win.addstr(*args) """ - ASCII_TRANSTABLE = str.maketrans( + ASCII_TRANSTABLE: ClassVar[dict[int, int]] = str.maketrans( '═─╴╒╤╕╪╘╧╛┌┬┐┼└┴┘│╞╡├┤▏▎▍▌▋▊▉█░▲▼␤' + GRAPH_SYMBOLS, '=--++++++++++++++||||||||||||||^v?' + '=' * len(GRAPH_SYMBOLS), ) - TERM_256COLOR = False + TERM_256COLOR: ClassVar[bool] = False - def __init__(self): - self.win = None # type: curses._CursesWindow - self.ascii = False + def __init__(self) -> None: + self.win: curses.window | None = None + self.no_unicode: bool = False - def addstr(self, *args, **kwargs): - if self.ascii: - args = [ + def addstr(self, *args: str | int | Color, **kwargs: str | int | Color) -> None: + if self.no_unicode: + args = [ # type: ignore[assignment] arg.translate(self.ASCII_TRANSTABLE) if isinstance(arg, str) else arg for arg in args ] + assert self.win is not None try: - self.win.addstr(*args, **kwargs) + self.win.addstr(*args, **kwargs) # type: ignore[arg-type] except curses.error: pass - def addnstr(self, *args, **kwargs): - if self.ascii: - args = [ + def addnstr(self, *args: str | int | Color, **kwargs: str | int | Color) -> None: + if self.no_unicode: + args = [ # type: ignore[assignment] arg.translate(self.ASCII_TRANSTABLE) if isinstance(arg, str) else arg for arg in args ] + assert self.win is not None try: - self.win.addnstr(*args, **kwargs) + self.win.addnstr(*args, **kwargs) # type: ignore[arg-type] except curses.error: pass - def addch(self, *args, **kwargs): - if self.ascii: - args = [ + def addch(self, *args: str | int | Color, **kwargs: str | int | Color) -> None: + if self.no_unicode: + args = [ # type: ignore[assignment] arg.translate(self.ASCII_TRANSTABLE) if isinstance(arg, str) else arg for arg in args ] + assert self.win is not None try: - self.win.addch(*args, **kwargs) + self.win.addch(*args, **kwargs) # type: ignore[arg-type] except curses.error: pass - def color(self, fg=-1, bg=-1, attr=0): + def color(self, fg: Color = -1, bg: Color = -1, attr: str | int = 0) -> int: """Change the colors from now on.""" return self.set_fg_bg_attr(fg, bg, attr) - def color_reset(self): + def color_reset(self) -> int: """Change the colors to the default colors.""" return self.color() - def color_at(self, y, x, width, *args, **kwargs): + def color_at( + self, + y: int, + x: int, + width: int, + *args: str | int | Color, + **kwargs: str | int | Color, + ) -> None: """Change the colors at the specified position.""" + assert self.win is not None try: - self.win.chgat(y, x, width, self.get_fg_bg_attr(*args, **kwargs)) + self.win.chgat(y, x, width, self.get_fg_bg_attr(*args, **kwargs)) # type: ignore[arg-type] except curses.error: pass @staticmethod - def get_fg_bg_attr(fg=-1, bg=-1, attr=0): + def get_fg_bg_attr(fg: Color = -1, bg: Color = -1, attr: str | int = 0) -> int: """Return the curses attribute for the given fg/bg/attr combination.""" if fg == -1 and bg == -1 and attr == 0: return BASE_ATTR @@ -268,7 +291,8 @@ class CursesShortcuts: return attr | BASE_ATTR return curses.color_pair(_get_color(fg, bg)) | attr | BASE_ATTR - def set_fg_bg_attr(self, fg=-1, bg=-1, attr=0): + def set_fg_bg_attr(self, fg: Color = -1, bg: Color = -1, attr: str | int = 0) -> int: + assert self.win is not None try: attr = self.get_fg_bg_attr(fg, bg, attr) self.win.attrset(attr) @@ -276,27 +300,28 @@ class CursesShortcuts: return 0 return attr - def update_size(self, termsize=None): + def update_size(self, termsize: tuple[int, int] | None = None) -> tuple[int, int]: if termsize is not None: return termsize self.update_lines_cols() + assert self.win is not None return self.win.getmaxyx() @staticmethod - def update_lines_cols(): + def update_lines_cols() -> None: curses.update_lines_cols() @staticmethod - def beep(): + def beep() -> None: curses.beep() @staticmethod - def flash(): + def flash() -> None: curses.flash() @staticmethod - def set_base_attr(attr=0): + def set_base_attr(attr: str | int = 0) -> None: global BASE_ATTR # pylint: disable=global-statement if isinstance(attr, str): diff --git a/nvitop/tui/library/messagebox.py b/nvitop/tui/library/messagebox.py index bf4b0b7..13d0c83 100644 --- a/nvitop/tui/library/messagebox.py +++ b/nvitop/tui/library/messagebox.py @@ -4,38 +4,80 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import curses import string import threading import time from functools import partial +from typing import TYPE_CHECKING, Literal from nvitop.tui.library import host from nvitop.tui.library.displayable import Displayable -from nvitop.tui.library.keybinding import normalize_keybinding +from nvitop.tui.library.keybinding import NAMED_SPECIAL_KEYS, normalize_keybinding from nvitop.tui.library.utils import cut_string from nvitop.tui.library.widestring import WideString -DIGITS = set(string.digits) +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + + from nvitop.tui.library.mouse import MouseEvent + from nvitop.tui.screens.base import BaseScreen, BaseSelectableScreen + from nvitop.tui.tui import TUI + + +__all__ = ['MessageBox'] + + +DIGITS: frozenset[str] = frozenset(string.digits) class MessageBox(Displayable): # pylint: disable=too-many-instance-attributes class Option: # pylint: disable=too-few-public-methods # pylint: disable-next=too-many-arguments - def __init__(self, name, key, callback, *, keys=(), attrs=()): - self.name = WideString(name) - self.offset = 0 - self.key = normalize_keybinding(key) - self.callback = callback - self.keys = tuple({normalize_keybinding(key) for key in keys}.difference({self.key})) - self.attrs = attrs + def __init__( + self, + name: str, + key: str, + callback: Callable[[], None] | None, + *, + keys: Iterable[str] = (), + attrs: tuple[dict[str, int | str], ...] = (), + ) -> None: + self.name: WideString = WideString(name) + self.offset: int = 0 + self.key: str = normalize_keybinding(key) + self.callback: Callable[[], None] | None = callback + self.keys: tuple[str, ...] = tuple( + set(map(normalize_keybinding, keys)).difference({self.key}), + ) + self.attrs: tuple[dict[str, int | str], ...] = attrs - def __str__(self): + def __call__(self) -> None: + if self.callback is not None: + self.callback() + + def __str__(self) -> str: return str(self.name) + root: TUI + parent: TUI + # pylint: disable-next=too-many-arguments - def __init__(self, message, options, *, default, yes, no, cancel, win, root): + def __init__( + self, + message: str, + options: list[MessageBox.Option], + *, + default: int | None, + yes: int | None, + no: int | None, + cancel: int, + win: curses.window, + root: TUI, + ) -> None: super().__init__(win, root) if default is None: @@ -43,8 +85,8 @@ class MessageBox(Displayable): # pylint: disable=too-many-instance-attributes if no is None: no = cancel - self.options = options - self.num_options = len(self.options) + self.options: list[MessageBox.Option] = options + self.num_options: int = len(self.options) assert cancel is not None assert self.num_options >= 2 @@ -52,27 +94,27 @@ class MessageBox(Displayable): # pylint: disable=too-many-instance-attributes assert 0 <= cancel < self.num_options assert 0 <= default < self.num_options - self.previous_focused = None - self.message = message - self.previous_keymap = root.keymaps.used_keymap - self.current = default - self.yes = yes - self.cancel = cancel - self.no = no # pylint: disable=invalid-name - self.timestamp = time.monotonic() + self.previous_focused: BaseScreen | None = None + self.message: str = message + self.previous_keymap: str = root.keymaps.used_keymap # type: ignore[assignment] + self.current: int = default + self.yes: int | None = yes + self.cancel: int = cancel + self.no: int = no # pylint: disable=invalid-name + self.timestamp: float = time.monotonic() - self.name_len = max(8, *(len(option.name) for option in options)) + self.name_len: int = max(8, *(len(option.name) for option in options)) for option in self.options: option.offset = (self.name_len - len(option.name)) // 2 option.name = option.name.center(self.name_len) - self.xy_mouse = None + self.xy_mouse: tuple[int, int] | None = None self.x, self.y = root.x, root.y - self.width = (self.name_len + 6) * self.num_options + 6 + self.width: int = (self.name_len + 6) * self.num_options + 6 self.init_keybindings() - lines = [] + lines: list[str | WideString] = [] for msg in self.message.splitlines(): words = iter(map(WideString, msg.split())) try: @@ -88,32 +130,33 @@ class MessageBox(Displayable): # pylint: disable=too-many-instance-attributes lines.append(word) if len(lines) == 1: lines[-1] = WideString(lines[-1]).center(self.width - 6) - lines = [f' │ {line.ljust(self.width - 6)} │ ' for line in lines] - lines = [ + raw_lines = [f' │ {line.ljust(self.width - 6)} │ ' for line in lines] + raw_lines = [ ' ╒' + '═' * (self.width - 4) + '╕ ', ' │' + ' ' * (self.width - 4) + '│ ', - *lines, + *raw_lines, ' │' + ' ' * (self.width - 4) + '│ ', ' │ ' + ' '.join(['┌' + '─' * (self.name_len + 2) + '┐'] * self.num_options) + ' │ ', ' │ ' + ' '.join(map('│ {} │'.format, self.options)) + ' │ ', ' │ ' + ' '.join(['└' + '─' * (self.name_len + 2) + '┘'] * self.num_options) + ' │ ', ' ╘' + '═' * (self.width - 4) + '╛ ', ] - self.lines = lines + self.lines: list[str] = raw_lines @property - def current(self): + def current(self) -> int: return self._current @current.setter - def current(self, value): + def current(self, value: int) -> None: self._current = value self.timestamp = time.monotonic() - def draw(self): + def draw(self) -> None: # pylint: disable=too-many-locals self.set_base_attr(attr=0) self.color_reset() + assert self.root.termsize is not None n_term_lines, n_term_cols = self.root.termsize height = len(self.lines) @@ -126,9 +169,10 @@ class MessageBox(Displayable): # pylint: disable=too-many-instance-attributes x_option_start = x_start + 6 + i * (self.name_len + 6) + option.offset for attr in option.attrs: attr = attr.copy() - y = y_option_start + attr.pop('y') - x = x_option_start + attr.pop('x') - self.color_at(y, x, **attr) + y = y_option_start + attr.pop('y') # type: ignore[operator] + x = x_option_start + attr.pop('x') # type: ignore[operator] + width: int = attr.pop('width') # type: ignore[assignment] + self.color_at(y, x, width, **attr) if self.xy_mouse is not None: x, y = self.xy_mouse @@ -152,22 +196,23 @@ class MessageBox(Displayable): # pylint: disable=too-many-instance-attributes ) for attr in option.attrs: attr = attr.copy() - y = y_option_start + attr.pop('y') - x = x_option_start + option.offset + attr.pop('x') + y = y_option_start + attr.pop('y') # type: ignore[operator] + x = x_option_start + option.offset + attr.pop('x') # type: ignore[operator] + width = attr.pop('width') # type: ignore[assignment] attr['fg'], attr['bg'] = attr.get('bg', -1), attr.get('fg', -1) attr['attr'] = self.get_fg_bg_attr(attr=attr.get('attr', 0)) - attr['attr'] |= self.get_fg_bg_attr(attr='standout | bold') - self.color_at(y, x, **attr) + attr['attr'] |= self.get_fg_bg_attr(attr='standout | bold') # type: ignore[operator] + self.color_at(y, x, width, **attr) - def finalize(self): + def finalize(self) -> None: self.xy_mouse = None super().finalize() - def press(self, key): + def press(self, key: int) -> bool: self.root.keymaps.use_keymap('messagebox') - self.root.press(key) + return self.root.press(key) - def click(self, event): + def click(self, event: MouseEvent) -> bool: if event.pressed(1) or event.pressed(3) or event.clicked(1) or event.clicked(3): self.xy_mouse = (event.x, event.y) return True @@ -176,7 +221,7 @@ class MessageBox(Displayable): # pylint: disable=too-many-instance-attributes self.current = (self.current + direction) % self.num_options return True - def apply(self, index=None, wait=None): + def apply(self, index: int | None = None, wait: bool | None = None) -> None: if index is None: index = self.current @@ -185,147 +230,147 @@ class MessageBox(Displayable): # pylint: disable=too-many-instance-attributes if (index != self.current and wait is None) or wait: self.current = index - def confirm(): + def confirm() -> None: time.sleep(0.25) curses.ungetch(curses.KEY_ENTER) threading.Thread(name='messagebox-confirm', target=confirm, daemon=True).start() return - callback = self.options[index].callback - if callback is not None: - callback() + option = self.options[index] + option() self.root.keymaps.clear_keymap('messagebox') self.root.keymaps.use_keymap(self.previous_keymap) self.root.need_redraw = True self.root.messagebox = None - def init_keybindings(self): # pylint: disable=too-many-branches - def select_previous(): + def init_keybindings(self) -> None: # pylint: disable=too-many-branches + def select_previous() -> None: self.current = (self.current - 1) % self.num_options - def select_next(): + def select_next() -> None: self.current = (self.current + 1) % self.num_options keymaps = self.root.keymaps - keymaps.clear_keymap('messagebox') + keymap = keymaps.clear_keymap('messagebox') for i, option in enumerate(self.options): keymaps.bind('messagebox', option.key, partial(self.apply, index=i)) for key in option.keys: - keymaps.copy('messagebox', option.key, key) + keymaps.alias('messagebox', option.key, key) - keymaps['messagebox'][keymaps.keybuffer.quantifier_key] = 'false' - if len(DIGITS.intersection(keymaps['messagebox'])) == 0 and self.num_options <= 9: + keymap[keymaps.keybuffer.quantifier_key] = keymaps.keybuffer.QUANTIFIER_KEY_FINISHED # type: ignore[assignment] + if len(DIGITS.intersection(keymap)) == 0 and self.num_options <= 9: for key_n, option in zip('123456789', self.options): - keymaps.copy('messagebox', option.key, key_n) + keymaps.alias('messagebox', option.key, key_n) - assert ( - len({'', '', '', ''}.intersection(keymaps['messagebox'])) == 0 + assert set(keymap).isdisjoint( + NAMED_SPECIAL_KEYS[key] for key in ('Enter', 'Esc', 'Left', 'Right') ) - if self.yes is not None and 'y' not in keymaps['messagebox']: - keymaps.copy('messagebox', self.options[self.yes].key, 'y') - if 'Y' not in keymaps['messagebox']: - keymaps.copy('messagebox', self.options[self.yes].key, 'Y') - if self.no is not None and 'n' not in keymaps['messagebox']: - keymaps.copy('messagebox', self.options[self.no].key, 'n') - if 'N' not in keymaps['messagebox']: - keymaps.copy('messagebox', self.options[self.no].key, 'N') + if self.yes is not None and ord('y') not in keymap: + keymaps.alias('messagebox', self.options[self.yes].key, 'y') + if ord('Y') not in keymap: + keymaps.alias('messagebox', self.options[self.yes].key, 'Y') + if self.no is not None and ord('n') not in keymap: + keymaps.alias('messagebox', self.options[self.no].key, 'n') + if ord('N') not in keymap: + keymaps.alias('messagebox', self.options[self.no].key, 'N') if self.cancel is not None: keymaps.bind('messagebox', '', partial(self.apply, index=self.cancel, wait=False)) - if 'q' not in keymaps['messagebox'] and 'Q' not in keymaps['messagebox']: - keymaps.copy('messagebox', '', 'q') - keymaps.copy('messagebox', '', 'Q') + if ord('q') not in keymap and ord('Q') not in keymap: + keymaps.alias('messagebox', '', 'q') + keymaps.alias('messagebox', '', 'Q') keymaps.bind('messagebox', '', self.apply) - if '' not in keymaps['messagebox']: - keymaps.copy('messagebox', '', '') + if NAMED_SPECIAL_KEYS['Space'] not in keymap: + keymaps.alias('messagebox', '', '') keymaps.bind('messagebox', '', select_previous) keymaps.bind('messagebox', '', select_next) - if ',' not in keymaps['messagebox'] and '.' not in keymaps['messagebox']: - keymaps.copy('messagebox', '', ',') - keymaps.copy('messagebox', '', '.') - if '<' not in keymaps['messagebox'] and '>' not in keymaps['messagebox']: - keymaps.copy('messagebox', '', '<') - keymaps.copy('messagebox', '', '>') - if '[' not in keymaps['messagebox'] and ']' not in keymaps['messagebox']: - keymaps.copy('messagebox', '', '[') - keymaps.copy('messagebox', '', ']') - if '' not in keymaps['messagebox'] and '' not in keymaps['messagebox']: - keymaps.copy('messagebox', '', '') - keymaps.copy('messagebox', '', '') + if ord(',') not in keymap and ord('.') not in keymap: + keymaps.alias('messagebox', '', ',') + keymaps.alias('messagebox', '', '.') + if ord('<') not in keymap and ord('>') not in keymap: + keymaps.alias('messagebox', '', '<') + keymaps.alias('messagebox', '', '>') + if ord('[') not in keymap and ord(']') not in keymap: + keymaps.alias('messagebox', '', '[') + keymaps.alias('messagebox', '', ']') + if NAMED_SPECIAL_KEYS['Tab'] not in keymap and NAMED_SPECIAL_KEYS['S-Tab'] not in keymap: + keymaps.alias('messagebox', '', '') + keymaps.alias('messagebox', '', '') + @staticmethod + def confirm_sending_signal_to_processes( + signal: Literal['terminate', 'kill', 'interrupt'], + screen: BaseSelectableScreen, + ) -> None: + assert signal in ('terminate', 'kill', 'interrupt') + default = {'terminate': 0, 'kill': 1, 'interrupt': 2}.get(signal) + processes = [] + for process in screen.selection.processes(): + try: + username = process.username() + except host.PsutilError: + username = 'N/A' + username = cut_string(username, maxlen=24, padstr='+') + processes.append(f'{process.pid}({username})') + if len(processes) == 0: + return + if len(processes) == 1: + message = f'Send signal to process {processes[0]}?' + else: + maxlen = max(map(len, processes)) + processes = [process.ljust(maxlen) for process in processes] + message = 'Send signal to the following processes?\n\n{}'.format(' '.join(processes)) -def send_signal(signal, panel): - assert signal in {'terminate', 'kill', 'interrupt'} - default = {'terminate': 0, 'kill': 1, 'interrupt': 2}.get(signal) - processes = [] - for process in panel.selection.processes(): - try: - username = process.username() - except host.PsutilError: - username = 'N/A' - username = cut_string(username, maxlen=24, padstr='+') - processes.append(f'{process.pid}({username})') - if len(processes) == 0: - return - if len(processes) == 1: - message = f'Send signal to process {processes[0]}?' - else: - maxlen = max(map(len, processes)) - processes = [process.ljust(maxlen) for process in processes] - message = 'Send signal to the following processes?\n\n{}'.format(' '.join(processes)) - - # pylint: disable=use-dict-literal - panel.root.messagebox = MessageBox( - message=message, - options=[ - MessageBox.Option( - 'SIGTERM', - 't', - panel.selection.terminate, - keys=('T',), - attrs=( - {'y': 0, 'x': 0, 'width': 7, 'fg': 'red'}, - {'y': 0, 'x': 3, 'width': 1, 'fg': 'red', 'attr': 'bold | underline'}, + screen.root.messagebox = MessageBox( + message=message, + options=[ + MessageBox.Option( + 'SIGTERM', + 't', + screen.selection.terminate, + keys=('T',), + attrs=( + {'y': 0, 'x': 0, 'width': 7, 'fg': 'red'}, + {'y': 0, 'x': 3, 'width': 1, 'fg': 'red', 'attr': 'bold | underline'}, + ), ), - ), - MessageBox.Option( - 'SIGKILL', - 'k', - panel.selection.kill, - keys=('K',), - attrs=( - {'y': 0, 'x': 0, 'width': 7, 'fg': 'red'}, - {'y': 0, 'x': 3, 'width': 1, 'fg': 'red', 'attr': 'bold | underline'}, + MessageBox.Option( + 'SIGKILL', + 'k', + screen.selection.kill, + keys=('K',), + attrs=( + {'y': 0, 'x': 0, 'width': 7, 'fg': 'red'}, + {'y': 0, 'x': 3, 'width': 1, 'fg': 'red', 'attr': 'bold | underline'}, + ), ), - ), - MessageBox.Option( - 'SIGINT', - 'i', - panel.selection.interrupt, - keys=('I',), - attrs=( - {'y': 0, 'x': 0, 'width': 6, 'fg': 'red'}, - {'y': 0, 'x': 3, 'width': 1, 'fg': 'red', 'attr': 'bold | underline'}, + MessageBox.Option( + 'SIGINT', + 'i', + screen.selection.interrupt, + keys=('I',), + attrs=( + {'y': 0, 'x': 0, 'width': 6, 'fg': 'red'}, + {'y': 0, 'x': 3, 'width': 1, 'fg': 'red', 'attr': 'bold | underline'}, + ), ), - ), - MessageBox.Option( - 'Cancel', - 'c', - None, - keys=('C',), - attrs=({'y': 0, 'x': 0, 'width': 1, 'attr': 'bold | underline'},), - ), - ], - default=default, - yes=None, - no=3, - cancel=3, - win=panel.win, - root=panel.root, - ) - # pylint: enable=use-dict-literal + MessageBox.Option( + 'Cancel', + 'c', + None, + keys=('C',), + attrs=({'y': 0, 'x': 0, 'width': 1, 'attr': 'bold | underline'},), + ), + ], + default=default, + yes=None, + no=3, + cancel=3, + win=screen.win, # type: ignore[arg-type] + root=screen.root, + ) diff --git a/nvitop/tui/library/mouse.py b/nvitop/tui/library/mouse.py index e6d6b13..e19d57c 100644 --- a/nvitop/tui/library/mouse.py +++ b/nvitop/tui/library/mouse.py @@ -4,41 +4,51 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import curses +from typing import TYPE_CHECKING, ClassVar + + +if TYPE_CHECKING: + from typing_extensions import Self # Python 3.11+ + + +__all__ = ['MouseEvent'] class MouseEvent: - PRESSED = [ + PRESSED: ClassVar[tuple[int, ...]] = ( 0, curses.BUTTON1_PRESSED, curses.BUTTON2_PRESSED, curses.BUTTON3_PRESSED, curses.BUTTON4_PRESSED, - ] - RELEASED = [ + ) + RELEASED: ClassVar[tuple[int, ...]] = ( 0, curses.BUTTON1_RELEASED, curses.BUTTON2_RELEASED, curses.BUTTON3_RELEASED, curses.BUTTON4_RELEASED, - ] - CLICKED = [ + ) + CLICKED: ClassVar[tuple[int, ...]] = ( 0, curses.BUTTON1_CLICKED, curses.BUTTON2_CLICKED, curses.BUTTON3_CLICKED, curses.BUTTON4_CLICKED, - ] - DOUBLE_CLICKED = [ + ) + DOUBLE_CLICKED: ClassVar[tuple[int, ...]] = ( 0, curses.BUTTON1_DOUBLE_CLICKED, curses.BUTTON2_DOUBLE_CLICKED, curses.BUTTON3_DOUBLE_CLICKED, curses.BUTTON4_DOUBLE_CLICKED, - ] - CTRL_SCROLLWHEEL_MULTIPLIER = 5 + ) + CTRL_SCROLLWHEEL_MULTIPLIER: ClassVar[int] = 5 - def __init__(self, state): + def __init__(self, state: tuple[int, int, int, int, int]) -> None: """Create a MouseEvent object from the result of win.getmouse().""" _, self.x, self.y, _, self.bstate = state @@ -46,39 +56,43 @@ class MouseEvent: # it's sufficient to add 0xFF to fix that error. if self.x < 0: self.x += 0xFF - if self.y < 0: self.y += 0xFF - def pressed(self, n): + @classmethod + def get(cls) -> Self: + """Get the mouse event.""" + return cls(curses.getmouse()) + + def pressed(self, n: int) -> bool: """Return whether the mouse key n is pressed.""" try: - return (self.bstate & MouseEvent.PRESSED[n]) != 0 + return bool(self.bstate & self.PRESSED[n]) except IndexError: return False - def released(self, n): + def released(self, n: int) -> bool: """Return whether the mouse key n is released.""" try: - return (self.bstate & MouseEvent.RELEASED[n]) != 0 + return bool(self.bstate & self.RELEASED[n]) except IndexError: return False - def clicked(self, n): + def clicked(self, n: int) -> bool: """Return whether the mouse key n is clicked.""" try: - return (self.bstate & MouseEvent.CLICKED[n]) != 0 + return bool(self.bstate & self.CLICKED[n]) except IndexError: return False - def double_clicked(self, n): + def double_clicked(self, n: int) -> bool: """Return whether the mouse key n is double clicked.""" try: - return (self.bstate & MouseEvent.DOUBLE_CLICKED[n]) != 0 + return bool(self.bstate & self.DOUBLE_CLICKED[n]) except IndexError: return False - def wheel_direction(self): + def wheel_direction(self) -> int: """Return the direction of the scroll action, 0 if there was none.""" # If the bstate > ALL_MOUSE_EVENTS, it's an invalid mouse button. # I interpret invalid buttons as "scroll down" because all tested @@ -91,14 +105,14 @@ class MouseEvent: return self.CTRL_SCROLLWHEEL_MULTIPLIER if self.ctrl() else 1 return 0 - def ctrl(self): - return self.bstate & curses.BUTTON_CTRL + def ctrl(self) -> bool: + return bool(self.bstate & curses.BUTTON_CTRL) - def alt(self): - return self.bstate & curses.BUTTON_ALT + def alt(self) -> bool: + return bool(self.bstate & curses.BUTTON_ALT) - def shift(self): - return self.bstate & curses.BUTTON_SHIFT + def shift(self) -> bool: + return bool(self.bstate & curses.BUTTON_SHIFT) - def key_invalid(self): + def key_invalid(self) -> bool: return self.bstate > curses.ALL_MOUSE_EVENTS diff --git a/nvitop/tui/library/process.py b/nvitop/tui/library/process.py index 1f26820..49af3fe 100644 --- a/nvitop/tui/library/process.py +++ b/nvitop/tui/library/process.py @@ -3,16 +3,29 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + from nvitop.api import NA, HostProcess, Snapshot, utilization2string from nvitop.api import GpuProcess as GpuProcessBase -from nvitop.api.host import WINDOWS, WSL +from nvitop.tui.library.utils import IS_WINDOWS, IS_WSL -__all__ = ['HostProcess', 'GpuProcess'] +if TYPE_CHECKING: + from typing_extensions import Self # Python 3.11+ + + from nvitop.tui.library.device import Device + + +__all__ = ['GpuProcess', 'HostProcess'] class GpuProcess(GpuProcessBase): - def __new__(cls, *args, **kwargs): + _snapshot: Snapshot | None + device: Device + + def __new__(cls, *args: Any, **kwargs: Any) -> Self: instance = super().__new__(cls, *args, **kwargs) instance._snapshot = None return instance @@ -21,7 +34,7 @@ class GpuProcess(GpuProcessBase): def snapshot(self) -> Snapshot: if self._snapshot is None: self.as_snapshot() - return self._snapshot + return self._snapshot # type: ignore[return-value] def host_snapshot(self) -> Snapshot: host_snapshot = super().host_snapshot() @@ -42,11 +55,15 @@ class GpuProcess(GpuProcessBase): return host_snapshot - def as_snapshot(self, *, host_process_snapshot_cache=None) -> Snapshot: + def as_snapshot( + self, + *, + host_process_snapshot_cache: dict[int, Snapshot] | None = None, + ) -> Snapshot: snapshot = super().as_snapshot(host_process_snapshot_cache=host_process_snapshot_cache) snapshot.type = snapshot.type.replace('C+G', 'X') - if snapshot.gpu_memory_human is NA and (WINDOWS or WSL): + if snapshot.gpu_memory_human is NA and (IS_WINDOWS or IS_WSL): snapshot.gpu_memory_human = 'WDDM:N/A' snapshot.cpu_percent_string = snapshot.host.cpu_percent_string diff --git a/nvitop/tui/library/selection.py b/nvitop/tui/library/selection.py index 0f934ec..50bf1fc 100644 --- a/nvitop/tui/library/selection.py +++ b/nvitop/tui/library/selection.py @@ -3,63 +3,76 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import signal import time +from typing import TYPE_CHECKING from weakref import WeakValueDictionary from nvitop.api import NA, Snapshot from nvitop.tui.library import host -from nvitop.tui.library.utils import LARGE_INTEGER, SUPERUSER, USERNAME +from nvitop.tui.library.utils import IS_SUPERUSER, IS_WINDOWS, LARGE_INTEGER, USERNAME + + +if TYPE_CHECKING: + from collections.abc import Callable + + from nvitop.tui.library.displayable import Displayable + from nvitop.tui.library.process import GpuProcess, HostProcess + + +__all__ = ['Selection'] class Selection: # pylint: disable=too-many-instance-attributes - def __init__(self, panel): - self.tagged = WeakValueDictionary() - self.panel = panel - self.index = None - self.within_window = True - self._process = None - self._username = None - self._ident = None + def __init__(self, displayable: Displayable) -> None: + self.tagged: WeakValueDictionary[int, GpuProcess | HostProcess] = WeakValueDictionary() + self.displayable: Displayable = displayable + self.index: int | None = None + self.within_window: bool = True + self._process: GpuProcess | HostProcess | None = None + self._username: str | None = None + self._ident: tuple | None = None @property - def identity(self): + def identity(self) -> tuple: if self._ident is None: - self._ident = self.process._ident # pylint: disable=protected-access + self._ident = self.process._ident # type: ignore[union-attr] # pylint: disable=protected-access return self._ident @property - def process(self): + def process(self) -> GpuProcess | HostProcess | None: return self._process @process.setter - def process(self, process): + def process(self, process: Snapshot | GpuProcess | HostProcess | None) -> None: if isinstance(process, Snapshot): process = process.real - self._process = process + self._process = process # type: ignore[assignment] self._ident = None @property - def pid(self): + def pid(self) -> int | None: try: return self.identity[0] except TypeError: return None @property - def username(self): + def username(self) -> str | None: if self._username is None: try: - self._username = self.process.username() + self._username = self.process.username() # type: ignore[union-attr] except host.PsutilError: self._username = NA return self._username - def move(self, direction=0): + def move(self, direction: int = 0) -> None: if direction == 0: return - processes = self.panel.snapshots + processes = self.displayable.snapshots # type: ignore[attr-defined] old_index = self.index if len(processes) > 0: @@ -69,7 +82,7 @@ class Selection: # pylint: disable=too-many-instance-attributes else: self.index = len(processes) - 1 if direction > 0 else 0 else: - self.index = min(max(0, self.index + direction), len(processes) - 1) + self.index = min(max(0, self.index + direction), len(processes) - 1) # type: ignore[operator] self.process = processes[self.index] if old_index is not None: @@ -77,33 +90,31 @@ class Selection: # pylint: disable=too-many-instance-attributes else: direction = 0 - if direction != 0 and self.panel.NAME == 'process': - self.panel.parent.move(direction) + if direction != 0 and self.displayable.NAME == 'process': # type: ignore[attr-defined] + self.displayable.parent.move(direction) # type: ignore[union-attr] else: self.clear() - def owned(self): + def owned(self) -> bool: if not self.is_set(): return False - if SUPERUSER: - return True - return self.username == USERNAME + return IS_SUPERUSER or self.username == USERNAME - def tag(self): + def tag(self) -> None: if self.is_set(): try: - del self.tagged[self.pid] + del self.tagged[self.pid] # type: ignore[arg-type] except KeyError: - self.tagged[self.pid] = self.process + self.tagged[self.pid] = self.process # type: ignore[index,assignment] - def processes(self): + def processes(self) -> tuple[GpuProcess | HostProcess, ...]: if len(self.tagged) > 0: return tuple(sorted(self.tagged.values(), key=lambda p: p.pid)) if self.owned() and self.within_window: - return (self.process,) + return (self.process,) # type: ignore[return-value] return () - def foreach(self, func): + def foreach(self, func: Callable[[GpuProcess | HostProcess], None]) -> None: flag = False for process in self.processes(): try: @@ -117,39 +128,40 @@ class Selection: # pylint: disable=too-many-instance-attributes time.sleep(0.25) self.clear() - def send_signal(self, sig): + def send_signal(self, sig: int) -> None: self.foreach(lambda process: process.send_signal(sig)) - def interrupt(self): + def interrupt(self) -> None: try: # pylint: disable-next=no-member - self.send_signal(signal.SIGINT if not host.WINDOWS else signal.CTRL_C_EVENT) + self.send_signal(signal.SIGINT if not IS_WINDOWS else signal.CTRL_C_EVENT) # type: ignore[attr-defined] except SystemError: pass - def terminate(self): + def terminate(self) -> None: self.foreach(lambda process: process.terminate()) - def kill(self): + def kill(self) -> None: self.foreach(lambda process: process.kill()) - def reset(self): + def reset(self) -> None: self.index = None self.within_window = True self._process = None self._username = None self._ident = None - def clear(self): + def clear(self) -> None: self.tagged.clear() self.reset() - def is_set(self): + def is_set(self) -> bool: return self.process is not None - __bool__ = is_set + def is_same(self, process: Snapshot | GpuProcess | HostProcess) -> bool: + if isinstance(process, Snapshot): + process = process.real - def is_same(self, process): try: return self.identity == process._ident # pylint: disable=protected-access except (AttributeError, TypeError): @@ -157,9 +169,10 @@ class Selection: # pylint: disable=too-many-instance-attributes return False - __eq__ = is_same + def is_same_on_host(self, process: Snapshot | GpuProcess | HostProcess) -> bool: + if isinstance(process, Snapshot): + process = process.real - def is_same_on_host(self, process): try: return self.identity[:2] == process._ident[:2] # pylint: disable=protected-access except (AttributeError, TypeError): @@ -167,5 +180,5 @@ class Selection: # pylint: disable=too-many-instance-attributes return False - def is_tagged(self, process): + def is_tagged(self, process: Snapshot | GpuProcess | HostProcess) -> bool: return process.pid in self.tagged diff --git a/nvitop/tui/library/utils.py b/nvitop/tui/library/utils.py index f7ad973..dfb38b3 100644 --- a/nvitop/tui/library/utils.py +++ b/nvitop/tui/library/utils.py @@ -3,13 +3,17 @@ # pylint: disable=missing-module-docstring,missing-function-docstring +from __future__ import annotations + import contextlib import math import os +from typing import Literal from nvitop.api import ( NA, GiB, + NaType, Snapshot, bytes2human, colored, @@ -17,18 +21,24 @@ from nvitop.api import ( timedelta2human, ttl_cache, ) -from nvitop.tui.library.host import WINDOWS, WSL, getuser, hostname +from nvitop.api.host import WINDOWS as IS_WINDOWS +from nvitop.api.host import WINDOWS_SUBSYSTEM_FOR_LINUX +from nvitop.tui.library.host import getuser, hostname from nvitop.tui.library.widestring import WideString __all__ = [ 'HOSTNAME', + 'IS_SUPERUSER', + 'IS_WINDOWS', + 'IS_WINDOWS_SUBSYSTEM_FOR_LINUX', + 'IS_WSL', 'LARGE_INTEGER', 'NA', - 'SUPERUSER', - 'USER_CONTEXT', 'USERNAME', + 'USER_CONTEXT', 'GiB', + 'NaType', 'Snapshot', 'bytes2human', 'colored', @@ -40,32 +50,38 @@ __all__ = [ ] -USERNAME = getuser() +USERNAME: str = getuser() -SUPERUSER = False +IS_SUPERUSER: bool = False with contextlib.suppress(AttributeError, OSError): - if WINDOWS: + if IS_WINDOWS: import ctypes - SUPERUSER = bool(ctypes.windll.shell32.IsUserAnAdmin()) + IS_SUPERUSER = bool(ctypes.windll.shell32.IsUserAnAdmin()) # type: ignore[attr-defined] else: try: - SUPERUSER = os.geteuid() == 0 + IS_SUPERUSER = os.geteuid() == 0 except AttributeError: - SUPERUSER = os.getuid() == 0 + IS_SUPERUSER = os.getuid() == 0 -HOSTNAME = hostname() -if WSL: +HOSTNAME: str = hostname() +IS_WINDOWS_SUBSYSTEM_FOR_LINUX = IS_WSL = bool(WINDOWS_SUBSYSTEM_FOR_LINUX) +if IS_WSL: HOSTNAME = f'{HOSTNAME} (WSL)' -USER_CONTEXT = f'{USERNAME}@{HOSTNAME}' +USER_CONTEXT: str = f'{USERNAME}@{HOSTNAME}' -LARGE_INTEGER = 65536 +LARGE_INTEGER: int = 65536 -def cut_string(s, maxlen, padstr='...', align='left'): - assert align in {'left', 'right'} +def cut_string( + s: object, + maxlen: int, + padstr: str = '...', + align: Literal['left', 'right'] = 'left', +) -> str: + assert align in ('left', 'right') if not isinstance(s, str): s = str(s) @@ -82,7 +98,7 @@ def cut_string(s, maxlen, padstr='...', align='left'): # pylint: disable=disallowed-name -def make_bar(prefix, percent, width, *, extra_text=''): +def make_bar(prefix: str, percent: float | str, width: int, *, extra_text: str = '') -> str: bar = f'{prefix}: ' if percent != NA and not (isinstance(percent, float) and not math.isfinite(percent)): if isinstance(percent, str) and percent.endswith('%'): @@ -96,7 +112,7 @@ def make_bar(prefix, percent, width, *, extra_text=''): if isinstance(percent, float) and len(f'{bar} {percent:.1f}%') <= width: text = f'{percent:.1f}%' else: - text = f'{min(round(percent), 100):d}%'.replace('100%', 'MAX') + text = f'{min(round(percent), 100):d}%'.replace('100%', 'MAX') # type: ignore[arg-type] else: bar += '░' * (width - len(bar) - 4) text = 'N/A' diff --git a/nvitop/tui/library/widestring.py b/nvitop/tui/library/widestring.py index 67600e0..9ad32a1 100644 --- a/nvitop/tui/library/widestring.py +++ b/nvitop/tui/library/widestring.py @@ -4,23 +4,33 @@ # pylint: disable=missing-module-docstring,missing-class-docstring +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal from unicodedata import east_asian_width -ASCIIONLY = set(map(chr, range(1, 128))) -NARROW = 1 -WIDE = 2 -WIDE_SYMBOLS = set('WF') +if TYPE_CHECKING: + from typing_extensions import Self # Python 3.11+ -def utf_char_width(string): +__all__ = ['WideString', 'wcslen'] + + +ASCIIONLY: frozenset[str] = frozenset(map(chr, range(1, 128))) +NARROW: Literal[1] = 1 +WIDE: Literal[2] = 2 +WIDE_SYMBOLS: frozenset[str] = frozenset('WF') + + +def utf_char_width(string: str) -> Literal[1, 2]: """Return the width of a single character.""" if east_asian_width(string) in WIDE_SYMBOLS: return WIDE return NARROW -def string_to_charlist(string): +def string_to_charlist(string: str) -> list[str]: """Return a list of characters with extra empty strings after wide chars.""" if ASCIIONLY.issuperset(string): return list(string) @@ -32,26 +42,26 @@ def string_to_charlist(string): return result -def wcslen(string): - """Return the length of a string with wide chars.""" +def wcslen(string: str | WideString) -> int: + # pylint: disable=wrong-spelling-in-docstring + """Return the length of a string with wide chars. + + >>> wcslen('poo') + 3 + >>> wcslen('十百千万') + 8 + >>> wcslen('a十') + 3 + """ return len(WideString(string)) -class WideString: # pylint: disable=too-few-public-methods,wrong-spelling-in-docstring - def __init__(self, string='', chars=None): - if isinstance(string, WideString): - string = string.string +class WideString: # pylint: disable=wrong-spelling-in-docstring + def __init__(self, string: str | WideString = '', chars: list[str] | None = None) -> None: + self.string: str = str(string) + self.chars: list[str] = string_to_charlist(self.string) if chars is None else chars - try: - self.string = str(string) - except UnicodeEncodeError: - self.string = string.encode('latin-1', 'ignore') - if chars is None: - self.chars = string_to_charlist(string) - else: - self.chars = chars - - def __add__(self, other): + def __add__(self, other: object) -> WideString: """ >>> (WideString('a') + WideString('b')).string 'ab' @@ -66,7 +76,7 @@ class WideString: # pylint: disable=too-few-public-methods,wrong-spelling-in-do return WideString(self.string + other.string, self.chars + other.chars) return NotImplemented - def __radd__(self, other): + def __radd__(self, other: object) -> WideString: """ >>> ('bc' + WideString('afd')).chars ['b', 'c', 'a', 'f', 'd'] @@ -77,27 +87,27 @@ class WideString: # pylint: disable=too-few-public-methods,wrong-spelling-in-do return WideString(other.string + self.string, other.chars + self.chars) return NotImplemented - def __iadd__(self, other): + def __iadd__(self, other: object) -> Self: new = self + other self.string = new.string self.chars = new.chars return self - def __str__(self): + def __str__(self) -> str: return self.string - def __repr__(self): + def __repr__(self) -> str: return f'<{self.__class__.__name__} {self.string!r}>' - def __eq__(self, other): + def __eq__(self, other: object) -> bool: if not isinstance(other, (str, WideString)): raise TypeError return str(self) == str(other) - def __hash__(self): + def __hash__(self) -> int: return hash(self.string) - def __getitem__(self, item): + def __getitem__(self, item: int | slice) -> WideString: """ >>> WideString('asdf')[2] @@ -109,21 +119,21 @@ class WideString: # pylint: disable=too-few-public-methods,wrong-spelling-in-do >>> WideString('asdf')[1:-100] - >>> WideString('モヒカン')[2:4] - - >>> WideString('モヒカン')[2:5] - - >>> WideString('モabカン')[2:5] + >>> WideString('十百千万')[2:4] + + >>> WideString('十百千万')[2:5] + + >>> WideString('十ab千万')[2:5] - >>> WideString('モヒカン')[1:5] - - >>> WideString('モヒカン')[:] - - >>> WideString('aモ')[0:3] - - >>> WideString('aモ')[0:2] + >>> WideString('十百千万')[1:5] + + >>> WideString('十百千万')[:] + + >>> WideString('a十')[0:3] + + >>> WideString('a十')[0:2] - >>> WideString('aモ')[0:1] + >>> WideString('a十')[0:1] """ if isinstance(item, slice): @@ -153,49 +163,49 @@ class WideString: # pylint: disable=too-few-public-methods,wrong-spelling-in-do return WideString(' ' + ''.join(self.chars[start : stop - 1])) return WideString(''.join(self.chars[start:stop])) - def __len__(self): + def __len__(self) -> int: """ >>> len(WideString('poo')) 3 - >>> len(WideString('モヒカン')) + >>> len(WideString('十百千万')) 8 """ return len(self.chars) - def ljust(self, width, fillchar=' '): + def ljust(self, width: int, fillchar: str = ' ') -> WideString: """ >>> WideString('poo').ljust(2) >>> WideString('poo').ljust(5) - >>> WideString('モヒカン').ljust(10) - + >>> WideString('十百千万').ljust(10) + """ if width > len(self): return WideString(self.string + fillchar * width)[:width] return self - def rjust(self, width, fillchar=' '): + def rjust(self, width: int, fillchar: str = ' ') -> WideString: """ >>> WideString('poo').rjust(2) >>> WideString('poo').rjust(5) - >>> WideString('モヒカン').rljust(10) - + >>> WideString('十百千万').rjust(10) + """ if width > len(self): return WideString(fillchar * width + self.string)[-width:] return self - def center(self, width, fillchar=' '): + def center(self, width: int, fillchar: str = ' ') -> WideString: """ >>> WideString('poo').center(2) >>> WideString('poo').center(5) - >>> WideString('モヒカン').center(10) - + >>> WideString('十百千万').center(10) + """ if width > len(self): left_width = (width - len(self)) // 2 @@ -203,29 +213,35 @@ class WideString: # pylint: disable=too-few-public-methods,wrong-spelling-in-do return WideString(fillchar * left_width + self.string + fillchar * right_width)[:width] return self - def strip(self, chars=None): + def strip(self, chars: str | None = None) -> WideString: """ >>> WideString(' poo ').strip() - >>> WideString(' モヒカン ').strip() - + >>> WideString(' 十百千万 ').strip() + """ return WideString(self.string.strip(chars)) - def lstrip(self, chars=None): + def lstrip(self, chars: str | None = None) -> WideString: """ >>> WideString(' poo ').lstrip() - >>> WideString(' モヒカン ').lstrip() - + >>> WideString(' 十百千万 ').lstrip() + """ return WideString(self.string.lstrip(chars)) - def rstrip(self, chars=None): + def rstrip(self, chars: str | None = None) -> WideString: """ >>> WideString(' poo ').rstrip() - >>> WideString(' モヒカン ').rstrip() - + >>> WideString(' 十百千万 ').rstrip() + """ return WideString(self.string.rstrip(chars)) + + +if __name__ == '__main__': + import doctest + + doctest.testmod() diff --git a/nvitop/tui/screens/__init__.py b/nvitop/tui/screens/__init__.py index 11812d2..666c6b4 100644 --- a/nvitop/tui/screens/__init__.py +++ b/nvitop/tui/screens/__init__.py @@ -3,8 +3,21 @@ # pylint: disable=missing-module-docstring +from nvitop.tui.screens.base import BaseScreen, BaseSelectableScreen from nvitop.tui.screens.environ import EnvironScreen from nvitop.tui.screens.help import HelpScreen from nvitop.tui.screens.main import BreakLoop, MainScreen from nvitop.tui.screens.metrics import ProcessMetricsScreen from nvitop.tui.screens.treeview import TreeViewScreen + + +__all__ = [ + 'BaseScreen', + 'BaseSelectableScreen', + 'BreakLoop', + 'EnvironScreen', + 'HelpScreen', + 'MainScreen', + 'ProcessMetricsScreen', + 'TreeViewScreen', +] diff --git a/nvitop/tui/screens/base.py b/nvitop/tui/screens/base.py new file mode 100644 index 0000000..fb4b368 --- /dev/null +++ b/nvitop/tui/screens/base.py @@ -0,0 +1,32 @@ +# This file is part of nvitop, the interactive NVIDIA-GPU process viewer. +# License: GNU GPL version 3. + +# pylint: disable=missing-module-docstring + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from nvitop.tui.library import DisplayableContainer, Selection + + +if TYPE_CHECKING: + from nvitop.tui.tui import TUI + + +__all__ = ['BaseScreen', 'BaseSelectableScreen'] + + +class BaseScreen(DisplayableContainer): + """Base class for all screens.""" + + root: TUI + parent: TUI + + NAME: ClassVar[str] + + +class BaseSelectableScreen(BaseScreen): + """Base class for all selectable screens.""" + + selection: Selection diff --git a/nvitop/tui/screens/environ.py b/nvitop/tui/screens/environ.py index 54c1155..324f596 100644 --- a/nvitop/tui/screens/environ.py +++ b/nvitop/tui/screens/environ.py @@ -3,42 +3,55 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + from collections import OrderedDict from functools import partial from itertools import islice +from typing import TYPE_CHECKING, ClassVar -from nvitop.tui.library import Displayable, GpuProcess, HostProcess, WideString, host +from nvitop.tui.library import GpuProcess, HostProcess, MouseEvent, WideString, host +from nvitop.tui.screens.base import BaseScreen -class EnvironScreen(Displayable): # pylint: disable=too-many-instance-attributes - NAME = 'environ' +if TYPE_CHECKING: + import curses - def __init__(self, win, root): + from nvitop.tui.tui import TUI + + +__all__ = ['EnvironScreen'] + + +class EnvironScreen(BaseScreen): # pylint: disable=too-many-instance-attributes + NAME: ClassVar[str] = 'environ' + + def __init__(self, *, win: curses.window, root: TUI) -> None: super().__init__(win, root) - self.this = HostProcess() + self.this: HostProcess = HostProcess() - self._process = None - self._environ = None - self.items = None - self.username = None - self.command = None + self._process: GpuProcess | HostProcess = self.this + self._environ: OrderedDict[WideString, WideString] | None = None + self.items: list[tuple[WideString, WideString]] | None = None + self.username: WideString = WideString('N/A') + self.command: WideString = WideString('N/A') - self.x_offset = 0 - self._y_offset = 0 - self.scroll_offset = 0 - self.y_mouse = None + self.x_offset: int = 0 + self._y_offset: int = 0 + self.scroll_offset: int = 0 + self.y_mouse: int | None = None - self._height = 0 + self._height: int = 0 self.x, self.y = root.x, root.y self.width, self.height = root.width, root.height @property - def process(self): + def process(self) -> GpuProcess | HostProcess: return self._process @process.setter - def process(self, value): + def process(self, value: GpuProcess | HostProcess | None) -> None: if value is None: value = self.this @@ -51,48 +64,47 @@ class EnvironScreen(Displayable): # pylint: disable=too-many-instance-attribute self.environ = None try: - self.command = self.process.command() + command = self.process.command() except host.PsutilError: - self.command = 'N/A' + command = 'N/A' try: - self.username = self.process.username() + username = self.process.username() except host.PsutilError: - self.username = 'N/A' + username = 'N/A' - self.command = WideString(self.command) - self.username = WideString(self.username) + self.command = WideString(command) + self.username = WideString(username) @property - def environ(self): + def environ(self) -> OrderedDict[WideString, WideString] | None: return self._environ @environ.setter - def environ(self, value): - newline = '␤' if not self.root.ascii else '?' + def environ(self, value: OrderedDict[str, str] | None) -> None: + newline = '␤' if not self.root.no_unicode else '?' - def normalize(s): + def normalize(s: str) -> str: return s.replace('\n', newline) if value is not None: self.items = [ - (WideString(key), WideString(f'{key}={normalize(value[key])}')) - for key in sorted(value.keys()) + (WideString(k), WideString(f'{k}={normalize(v)}')) for k, v in sorted(value.items()) ] - value = OrderedDict(self.items) + self._environ = OrderedDict(self.items) else: self.items = None - self._environ = value + self._environ = None self.x_offset = 0 self.y_offset = 0 self.scroll_offset = 0 @property - def height(self): + def height(self) -> int: return self._height @height.setter - def height(self, value): + def height(self, value: int) -> None: self._height = value try: self.y_offset = self.y_offset @@ -100,15 +112,15 @@ class EnvironScreen(Displayable): # pylint: disable=too-many-instance-attribute pass @property - def display_height(self): + def display_height(self) -> int: return self.height - 2 @property - def y_offset(self): + def y_offset(self) -> int: return self._y_offset @y_offset.setter - def y_offset(self, value): + def y_offset(self, value: int) -> None: if self.environ is None: self._y_offset = 0 self.scroll_offset = 0 @@ -122,7 +134,7 @@ class EnvironScreen(Displayable): # pylint: disable=too-many-instance-attribute self.scroll_offset = self.y_offset - self.display_height + 1 self.scroll_offset = min(self.scroll_offset, self.y_offset) - def move(self, direction, wheel=False): + def move(self, direction: int, wheel: bool = False) -> None: if self.environ is not None and wheel: n_items = len(self.environ) old_scroll_offset = self.scroll_offset @@ -134,7 +146,7 @@ class EnvironScreen(Displayable): # pylint: disable=too-many-instance-attribute self._y_offset += self.scroll_offset - old_scroll_offset self.y_offset += direction - def update_size(self, termsize=None): + def update_size(self, termsize: tuple[int, int] | None = None) -> tuple[int, int]: n_term_lines, n_term_cols = termsize = super().update_size(termsize=termsize) self.width = n_term_cols - self.x @@ -142,7 +154,7 @@ class EnvironScreen(Displayable): # pylint: disable=too-many-instance-attribute return termsize - def draw(self): + def draw(self) -> None: self.color_reset() if isinstance(self.process, GpuProcess): @@ -168,11 +180,12 @@ class EnvironScreen(Displayable): # pylint: disable=too-many-instance-attribute self.color_at(self.y + 2, self.x, width=self.width, fg='cyan', attr='reverse') return + assert self.items is not None items = islice(self.items, self.scroll_offset, self.scroll_offset + self.display_height) for y, (key, line) in enumerate(items, start=self.y + 2): key_length = len(key) - line = str(line[self.x_offset :].ljust(self.width)[: self.width]) - self.addstr(y, self.x, line) + raw_line = str(line[self.x_offset :].ljust(self.width)[: self.width]) + self.addstr(y, self.x, raw_line) if self.x_offset < key_length: self.color_at(y, self.x, width=key_length - self.x_offset, fg='blue', attr='bold') if self.x_offset < key_length + 1: @@ -184,15 +197,15 @@ class EnvironScreen(Displayable): # pylint: disable=too-many-instance-attribute if y == self.y + 2 - self.scroll_offset + self.y_offset: self.color_at(y, self.x, width=self.width, fg='cyan', attr='bold | reverse') - def finalize(self): + def finalize(self) -> None: self.y_mouse = None super().finalize() - def press(self, key): + def press(self, key: int) -> bool: self.root.keymaps.use_keymap('environ') - self.root.press(key) + return self.root.press(key) - def click(self, event): + def click(self, event: MouseEvent) -> bool: if event.pressed(1) or event.pressed(3) or event.clicked(1) or event.clicked(3): self.y_mouse = event.y return True @@ -204,44 +217,44 @@ class EnvironScreen(Displayable): # pylint: disable=too-many-instance-attribute self.move(direction=direction, wheel=True) return True - def init_keybindings(self): - def refresh_environ(): - self.process = self.root.previous_screen.selection.process + def init_keybindings(self) -> None: + def refresh_environ() -> None: + self.process = self.root.previous_screen.selection.process # type: ignore[attr-defined] self.need_redraw = True - def environ_left(): + def environ_left() -> None: self.x_offset = max(0, self.x_offset - 5) - def environ_right(): + def environ_right() -> None: self.x_offset += 5 - def environ_begin(): + def environ_begin() -> None: self.x_offset = 0 - def environ_move(direction): + def environ_move(direction: int) -> None: self.move(direction=direction) keymaps = self.root.keymaps keymaps.bind('environ', 'r', refresh_environ) - keymaps.copy('environ', 'r', 'R') - keymaps.copy('environ', 'r', '') - keymaps.copy('environ', 'r', '') + keymaps.alias('environ', 'r', 'R') + keymaps.alias('environ', 'r', '') + keymaps.alias('environ', 'r', '') keymaps.bind('environ', '', environ_left) - keymaps.copy('environ', '', '') + keymaps.alias('environ', '', '') keymaps.bind('environ', '', environ_right) - keymaps.copy('environ', '', '') + keymaps.alias('environ', '', '') keymaps.bind('environ', '', environ_begin) - keymaps.copy('environ', '', '^') + keymaps.alias('environ', '', '^') keymaps.bind('environ', '', partial(environ_move, direction=-1)) - keymaps.copy('environ', '', '') - keymaps.copy('environ', '', '') - keymaps.copy('environ', '', '') - keymaps.copy('environ', '', '[') + keymaps.alias('environ', '', '') + keymaps.alias('environ', '', '') + keymaps.alias('environ', '', '') + keymaps.alias('environ', '', '[') keymaps.bind('environ', '', partial(environ_move, direction=+1)) - keymaps.copy('environ', '', '') - keymaps.copy('environ', '', '') - keymaps.copy('environ', '', '') - keymaps.copy('environ', '', ']') + keymaps.alias('environ', '', '') + keymaps.alias('environ', '', '') + keymaps.alias('environ', '', '') + keymaps.alias('environ', '', ']') keymaps.bind('environ', '', partial(environ_move, direction=-(1 << 20))) keymaps.bind('environ', '', partial(environ_move, direction=+(1 << 20))) diff --git a/nvitop/tui/screens/help.py b/nvitop/tui/screens/help.py index 0188de5..a4afc5f 100644 --- a/nvitop/tui/screens/help.py +++ b/nvitop/tui/screens/help.py @@ -3,10 +3,24 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring -from nvitop.tui.library import Device, Displayable, MouseEvent +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from nvitop.tui.library import Device, MouseEvent +from nvitop.tui.screens.base import BaseScreen from nvitop.version import __version__ +if TYPE_CHECKING: + import curses + + from nvitop.tui.tui import TUI + + +__all__ = ['HelpScreen'] + + HELP_TEMPLATE = r"""nvitop {} - (C) Xuehai Pan, 2021-2025. Released under the GNU GPLv3 License. @@ -42,13 +56,13 @@ Press any key to return. """ -class HelpScreen(Displayable): # pylint: disable=too-many-instance-attributes - NAME = 'help' +class HelpScreen(BaseScreen): # pylint: disable=too-many-instance-attributes + NAME: ClassVar[str] = 'help' - def __init__(self, win, root): + def __init__(self, *, win: curses.window, root: TUI) -> None: super().__init__(win, root) - self.infos = ( + self.infos: list[str] = ( HELP_TEMPLATE.format( __version__, *Device.GPU_UTILIZATION_THRESHOLDS, @@ -58,7 +72,7 @@ class HelpScreen(Displayable): # pylint: disable=too-many-instance-attributes .strip() .splitlines() ) - self.color_matrix = { + self.color_matrix: dict[int, tuple[str | None, str | None]] = { 9: ('green', 'green'), 10: ('green', 'green'), 12: ('cyan', 'yellow'), @@ -73,10 +87,10 @@ class HelpScreen(Displayable): # pylint: disable=too-many-instance-attributes } self.x, self.y = root.x, root.y - self.width = max(map(len, self.infos)) - self.height = len(self.infos) + self.width: int = max(map(len, self.infos)) + self.height: int = len(self.infos) - def draw(self): + def draw(self) -> None: if not self.need_redraw: return @@ -108,6 +122,6 @@ class HelpScreen(Displayable): # pylint: disable=too-many-instance-attributes if right is not None: self.color_at(self.y + dy, self.x + 39, width=13, fg=right, attr='bold') - def press(self, key): + def press(self, key: int) -> bool: self.root.keymaps.use_keymap('help') - self.root.press(key) + return self.root.press(key) diff --git a/nvitop/tui/screens/main/__init__.py b/nvitop/tui/screens/main/__init__.py index c2f0944..929b6c7 100644 --- a/nvitop/tui/screens/main/__init__.py +++ b/nvitop/tui/screens/main/__init__.py @@ -3,47 +3,86 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import threading from functools import partial +from typing import TYPE_CHECKING, ClassVar, NoReturn -from nvitop.tui.library import LARGE_INTEGER, DisplayableContainer, MouseEvent, send_signal -from nvitop.tui.screens.main.device import DevicePanel -from nvitop.tui.screens.main.host import HostPanel -from nvitop.tui.screens.main.process import ProcessPanel +from nvitop.tui.library import ( + LARGE_INTEGER, + Device, + Displayable, + MessageBox, + MouseEvent, + Selection, + Snapshot, +) +from nvitop.tui.screens.base import BaseScreen, BaseSelectableScreen +from nvitop.tui.screens.main.panels import DevicePanel, HostPanel, OrderName, ProcessPanel + + +if TYPE_CHECKING: + import curses + from collections.abc import Callable, Iterable + + from nvitop.tui.tui import TUI, MonitorMode + + +__all__ = ['BreakLoop', 'MainScreen'] class BreakLoop(Exception): # noqa: N818 pass -class MainScreen(DisplayableContainer): # pylint: disable=too-many-instance-attributes - NAME = 'main' +class MainScreen(BaseSelectableScreen): # pylint: disable=too-many-instance-attributes + NAME: ClassVar[str] = 'main' - # pylint: disable-next=redefined-builtin,too-many-arguments,too-many-locals,too-many-statements - def __init__(self, devices, filters, *, ascii, mode, win, root): + # pylint: disable-next=too-many-arguments,too-many-locals,too-many-statements + def __init__( + self, + devices: list[Device], + filters: Iterable[Callable[[Snapshot], bool]], + *, + no_unicode: bool, + mode: MonitorMode, + win: curses.window | None, + root: TUI, + ) -> None: super().__init__(win, root) - self.width = root.width + self.width: int = root.width assert mode in {'auto', 'full', 'compact'} - compact = mode == 'compact' - self.mode = mode - self._compact = compact + compact: bool = mode == 'compact' + self.mode: MonitorMode = mode + self._compact: bool = compact - self.devices = devices - self.device_count = len(self.devices) + self.devices: list[Device] = devices + self.device_count: int = len(self.devices) self.snapshot_lock = threading.Lock() - self.device_panel = DevicePanel(self.devices, compact, win=win, root=root) + self.device_panel: DevicePanel = DevicePanel( + self.devices, + compact, + win=win, + root=root, + ) self.device_panel.focused = False self.add_child(self.device_panel) - self.host_panel = HostPanel(self.device_panel.leaf_devices, compact, win=win, root=root) + self.host_panel: HostPanel = HostPanel( + self.device_panel.leaf_devices, + compact, + win=win, + root=root, + ) self.host_panel.focused = False self.add_child(self.host_panel) - self.process_panel = ProcessPanel( + self.process_panel: ProcessPanel = ProcessPanel( self.device_panel.leaf_devices, compact, filters, @@ -53,13 +92,13 @@ class MainScreen(DisplayableContainer): # pylint: disable=too-many-instance-att self.process_panel.focused = False self.add_child(self.process_panel) - self.selection = self.process_panel.selection + self.selection: Selection = self.process_panel.selection - self.ascii = ascii - self.device_panel.ascii = self.ascii - self.host_panel.ascii = self.ascii - self.process_panel.ascii = self.ascii - if ascii: + self.no_unicode: bool = no_unicode + self.device_panel.no_unicode = self.no_unicode + self.host_panel.no_unicode = self.no_unicode + self.process_panel.no_unicode = self.no_unicode + if no_unicode: self.host_panel.full_height = self.host_panel.height = self.host_panel.compact_height self.x, self.y = root.x, root.y @@ -70,16 +109,16 @@ class MainScreen(DisplayableContainer): # pylint: disable=too-many-instance-att self.height = self.device_panel.height + self.host_panel.height + self.process_panel.height @property - def compact(self): + def compact(self) -> bool: return self._compact @compact.setter - def compact(self, value): + def compact(self, value: bool) -> None: if self._compact != value: self.need_redraw = True self._compact = value - def update_size(self, termsize=None): + def update_size(self, termsize: tuple[int, int] | None = None) -> tuple[int, int]: n_term_lines, n_term_cols = termsize = super().update_size(termsize=termsize) self.width = n_term_cols - self.x @@ -126,7 +165,7 @@ class MainScreen(DisplayableContainer): # pylint: disable=too-many-instance-att return termsize - def move(self, direction=0): + def move(self, direction: int = 0) -> None: if direction == 0: return @@ -134,7 +173,7 @@ class MainScreen(DisplayableContainer): # pylint: disable=too-many-instance-att self.update_size() self.need_redraw = True - def poke(self): + def poke(self) -> None: super().poke() height = self.device_panel.height + self.host_panel.height + self.process_panel.height @@ -142,12 +181,12 @@ class MainScreen(DisplayableContainer): # pylint: disable=too-many-instance-att self.update_size() self.need_redraw = True - def draw(self): + def draw(self) -> None: self.color_reset() super().draw() - def print(self): + def print(self) -> None: if self.device_count > 0: print_width = min(panel.print_width() for panel in self.container) self.width = max(print_width, min(self.width, 100)) @@ -157,115 +196,139 @@ class MainScreen(DisplayableContainer): # pylint: disable=too-many-instance-att panel.width = self.width panel.print() - def __contains__(self, item): + def __contains__(self, item: Displayable | MouseEvent | tuple[int, int]) -> bool: if self.visible and isinstance(item, MouseEvent): return True return super().__contains__(item) - def init_keybindings(self): + def init_keybindings(self) -> None: # pylint: disable=too-many-locals,too-many-statements - def quit(): # pylint: disable=redefined-builtin + def quit() -> NoReturn: # pylint: disable=redefined-builtin raise BreakLoop - def change_mode(mode): + def change_mode(mode: MonitorMode) -> None: self.mode = mode self.root.update_size() - def force_refresh(): + def force_refresh() -> None: select_clear() host_begin() self.y = self.root.y self.root.update_size() self.root.need_redraw = True - def screen_move(direction): + def screen_move(direction: int) -> None: self.move(direction) - def host_left(): + def host_left() -> None: self.process_panel.host_offset -= 2 - def host_right(): + def host_right() -> None: self.process_panel.host_offset += 2 - def host_begin(): + def host_begin() -> None: self.process_panel.host_offset = -1 - def host_end(): + def host_end() -> None: self.process_panel.host_offset = LARGE_INTEGER - def select_move(direction): + def select_move(direction: int) -> None: self.selection.move(direction=direction) - def select_clear(): + def select_clear() -> None: self.selection.clear() - def tag(): + def tag() -> None: self.selection.tag() select_move(direction=+1) - def sort_by(order, reverse): + def sort_by(order: OrderName, reverse: bool) -> None: self.process_panel.order = order self.process_panel.reverse = reverse self.root.update_size() - def order_previous(): + def order_previous() -> None: sort_by(order=ProcessPanel.ORDERS[self.process_panel.order].previous, reverse=False) - def order_next(): + def order_next() -> None: sort_by(order=ProcessPanel.ORDERS[self.process_panel.order].next, reverse=False) - def order_reverse(): + def order_reverse() -> None: sort_by(order=self.process_panel.order, reverse=not self.process_panel.reverse) keymaps = self.root.keymaps keymaps.bind('main', 'q', quit) - keymaps.copy('main', 'q', 'Q') + keymaps.alias('main', 'q', 'Q') keymaps.bind('main', 'a', partial(change_mode, mode='auto')) keymaps.bind('main', 'f', partial(change_mode, mode='full')) keymaps.bind('main', 'c', partial(change_mode, mode='compact')) keymaps.bind('main', 'r', force_refresh) - keymaps.copy('main', 'r', 'R') - keymaps.copy('main', 'r', '') - keymaps.copy('main', 'r', '') + keymaps.alias('main', 'r', 'R') + keymaps.alias('main', 'r', '') + keymaps.alias('main', 'r', '') keymaps.bind('main', '', partial(screen_move, direction=-1)) - keymaps.copy('main', '', '[') - keymaps.copy('main', '', '') + keymaps.alias('main', '', '[') + keymaps.alias('main', '', '') keymaps.bind('main', '', partial(screen_move, direction=+1)) - keymaps.copy('main', '', ']') - keymaps.copy('main', '', '') + keymaps.alias('main', '', ']') + keymaps.alias('main', '', '') keymaps.bind('main', '', host_left) - keymaps.copy('main', '', '') + keymaps.alias('main', '', '') keymaps.bind('main', '', host_right) - keymaps.copy('main', '', '') + keymaps.alias('main', '', '') keymaps.bind('main', '', host_begin) - keymaps.copy('main', '', '^') + keymaps.alias('main', '', '^') keymaps.bind('main', '', host_end) - keymaps.copy('main', '', '$') + keymaps.alias('main', '', '$') keymaps.bind('main', '', partial(select_move, direction=-1)) - keymaps.copy('main', '', '') - keymaps.copy('main', '', '') + keymaps.alias('main', '', '') + keymaps.alias('main', '', '') keymaps.bind('main', '', partial(select_move, direction=+1)) - keymaps.copy('main', '', '') - keymaps.copy('main', '', '') + keymaps.alias('main', '', '') + keymaps.alias('main', '', '') keymaps.bind('main', '', partial(select_move, direction=-(1 << 20))) keymaps.bind('main', '', partial(select_move, direction=+(1 << 20))) keymaps.bind('main', '', select_clear) keymaps.bind('main', '', tag) - keymaps.bind('main', 'T', partial(send_signal, signal='terminate', panel=self)) - keymaps.bind('main', 'K', partial(send_signal, signal='kill', panel=self)) - keymaps.copy('main', 'K', 'k') - keymaps.bind('main', '', partial(send_signal, signal='interrupt', panel=self)) - keymaps.copy('main', '', 'I') + keymaps.bind( + 'main', + 'T', + partial( + MessageBox.confirm_sending_signal_to_processes, + signal='terminate', + screen=self, + ), + ) + keymaps.bind( + 'main', + 'K', + partial( + MessageBox.confirm_sending_signal_to_processes, + signal='kill', + screen=self, + ), + ) + keymaps.alias('main', 'K', 'k') + keymaps.bind( + 'main', + '', + partial( + MessageBox.confirm_sending_signal_to_processes, + signal='interrupt', + screen=self, + ), + ) + keymaps.alias('main', '', 'I') keymaps.bind('main', ',', order_previous) - keymaps.copy('main', ',', '<') + keymaps.alias('main', ',', '<') keymaps.bind('main', '.', order_next) - keymaps.copy('main', '.', '>') + keymaps.alias('main', '.', '>') keymaps.bind('main', '/', order_reverse) for name, order in ProcessPanel.ORDERS.items(): keymaps.bind( diff --git a/nvitop/tui/screens/main/panels/__init__.py b/nvitop/tui/screens/main/panels/__init__.py new file mode 100644 index 0000000..15e3afe --- /dev/null +++ b/nvitop/tui/screens/main/panels/__init__.py @@ -0,0 +1,19 @@ +# This file is part of nvitop, the interactive NVIDIA-GPU process viewer. +# License: GNU GPL version 3. + +# pylint: disable=missing-module-docstring + +from nvitop.tui.screens.main.panels.base import BasePanel, BaseSelectablePanel +from nvitop.tui.screens.main.panels.device import DevicePanel +from nvitop.tui.screens.main.panels.host import HostPanel +from nvitop.tui.screens.main.panels.process import OrderName, ProcessPanel + + +__all__ = [ + 'BasePanel', + 'BaseSelectablePanel', + 'DevicePanel', + 'HostPanel', + 'OrderName', + 'ProcessPanel', +] diff --git a/nvitop/tui/screens/main/panels/base.py b/nvitop/tui/screens/main/panels/base.py new file mode 100644 index 0000000..ea5e219 --- /dev/null +++ b/nvitop/tui/screens/main/panels/base.py @@ -0,0 +1,34 @@ +# This file is part of nvitop, the interactive NVIDIA-GPU process viewer. +# License: GNU GPL version 3. + +# pylint: disable=missing-module-docstring + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from nvitop.tui.library import Displayable, Selection + + +if TYPE_CHECKING: + from nvitop.tui.screens.main import MainScreen + from nvitop.tui.tui import TUI + + +__all__ = ['BasePanel', 'BaseSelectablePanel'] + + +class BasePanel(Displayable): + """Base class for all panels.""" + + root: TUI + parent: MainScreen + + NAME: ClassVar[str] + SNAPSHOT_INTERVAL: ClassVar[float] = 0.5 + + +class BaseSelectablePanel(BasePanel): + """Base class for all selectable panels.""" + + selection: Selection diff --git a/nvitop/tui/screens/main/device.py b/nvitop/tui/screens/main/panels/device.py similarity index 84% rename from nvitop/tui/screens/main/device.py rename to nvitop/tui/screens/main/panels/device.py index 363d44a..f2e0873 100644 --- a/nvitop/tui/screens/main/device.py +++ b/nvitop/tui/screens/main/panels/device.py @@ -3,36 +3,56 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import threading import time +from typing import TYPE_CHECKING, ClassVar from nvitop.tui.library import ( + IS_WINDOWS, NA, - WINDOWS, Device, - Displayable, + MigDevice, + Snapshot, colored, cut_string, make_bar, ttl_cache, ) +from nvitop.tui.screens.main.panels.base import BasePanel from nvitop.version import __version__ -class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes - NAME = 'device' - SNAPSHOT_INTERVAL = 0.5 +if TYPE_CHECKING: + import curses - def __init__(self, devices, compact, win, root): + from nvitop.tui.tui import TUI + + +__all__ = ['DevicePanel'] + + +class DevicePanel(BasePanel): # pylint: disable=too-many-instance-attributes + NAME: ClassVar[str] = 'device' + SNAPSHOT_INTERVAL: ClassVar[float] = 0.5 + + def __init__( + self, + devices: list[Device], + compact: bool, + win: curses.window | None, + root: TUI, + ) -> None: super().__init__(win, root) - self.devices = devices - self.device_count = len(self.devices) + self.devices: list[Device] = devices + self.device_count: int = len(self.devices) - self.all_devices = [] - self.leaf_devices = [] - self.mig_device_counts = [0] * self.device_count - self.mig_enabled_device_count = 0 + self.all_devices: list[Device | MigDevice] = [] + self.leaf_devices: list[Device | MigDevice] = [] + self.mig_device_counts: list[int] = [0] * self.device_count + self.mig_enabled_device_count: int = 0 for i, device in enumerate(self.devices): self.all_devices.append(device) @@ -45,27 +65,27 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes else: self.leaf_devices.append(device) - self.mig_device_count = sum(self.mig_device_counts) - self.all_device_count = len(self.all_devices) - self.leaf_device_count = len(self.leaf_devices) + self.mig_device_count: int = sum(self.mig_device_counts) + self.all_device_count: int = len(self.all_devices) + self.leaf_device_count: int = len(self.leaf_devices) - self._compact = compact - self.width = max(79, root.width) - self.compact_height = ( + self._compact: bool = compact + self.width: int = max(79, root.width) + self.compact_height: int = ( 4 + 2 * (self.device_count + 1) + self.mig_device_count + self.mig_enabled_device_count ) - self.full_height = self.compact_height + self.device_count + 1 - self.height = self.compact_height if compact else self.full_height + self.full_height: int = self.compact_height + self.device_count + 1 + self.height: int = self.compact_height if compact else self.full_height if self.device_count == 0: self.height = self.full_height = self.compact_height = 6 - self.driver_version = Device.driver_version() - self.cuda_driver_version = Device.cuda_driver_version() + self.driver_version: str = Device.driver_version() + self.cuda_driver_version: str = Device.cuda_driver_version() - self._snapshot_buffer = [] - self._snapshots = [] + self._snapshot_buffer: list[Snapshot] = [] + self._snapshots: list[Snapshot] = [] self.snapshot_lock = threading.Lock() - self.snapshots = self.take_snapshots() + self.snapshots: list[Snapshot] = self.take_snapshots() self._snapshot_daemon = threading.Thread( name='device-snapshot-daemon', target=self._snapshot_target, @@ -73,30 +93,30 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes ) self._daemon_running = threading.Event() - self.formats_compact = [ + self.formats_compact: list[str] = [ '│ {physical_index:>3} {fan_speed_string:>3} {temperature_string:>4} ' '{performance_state:>3} {power_status:>12} ' '│ {memory_usage:>20} │ {gpu_utilization_string:>7} {compute_mode:>11} │', ] - self.formats_full = [ + self.formats_full: list[str] = [ '│ {physical_index:>3} {name:<18} {persistence_mode:<4} ' '│ {bus_id:<16} {display_active:>3} │ {total_volatile_uncorrected_ecc_errors:>20} │', '│ {fan_speed_string:>3} {temperature_string:>4} {performance_state:>4} {power_status:>12} ' '│ {memory_usage:>20} │ {gpu_utilization_string:>7} {compute_mode:>11} │', ] - self.mig_formats = [ + self.mig_formats: list[str] = [ '│{physical_index:>2}:{mig_index:<2}{name:>12} @ GI/CI:{gpu_instance_id:>2}/{compute_instance_id:<2}' '│ {memory_usage:>20} │ BAR1: {bar1_memory_used_human:>8} / {bar1_memory_percent_string:>3} │', ] - if WINDOWS: + if IS_WINDOWS: self.formats_full[0] = self.formats_full[0].replace( 'persistence_mode', 'current_driver_model', ) - self.support_mig = any('N/A' not in device.mig_mode for device in self.snapshots) + self.support_mig: bool = any('N/A' not in device.mig_mode for device in self.snapshots) if self.support_mig: self.formats_full[0] = self.formats_full[0].replace( '{total_volatile_uncorrected_ecc_errors:>20}', @@ -104,48 +124,48 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes ) @property - def width(self): + def width(self) -> int: return self._width @width.setter - def width(self, value): + def width(self, value: int) -> None: width = max(79, value) if self._width != width and self.visible: self.need_redraw = True self._width = width @property - def compact(self): + def compact(self) -> bool: return self._compact @compact.setter - def compact(self, value): + def compact(self, value: bool) -> None: if self._compact != value: self.need_redraw = True self._compact = value self.height = self.compact_height if self.compact else self.full_height @property - def snapshots(self): + def snapshots(self) -> list[Snapshot]: return self._snapshots @snapshots.setter - def snapshots(self, snapshots): + def snapshots(self, snapshots: list[Snapshot]) -> None: with self.snapshot_lock: self._snapshots = snapshots @classmethod - def set_snapshot_interval(cls, interval): + def set_snapshot_interval(cls, interval: float) -> None: assert interval > 0.0 interval = float(interval) cls.SNAPSHOT_INTERVAL = min(interval / 3.0, 1.0) - cls.take_snapshots = ttl_cache(ttl=interval)( - cls.take_snapshots.__wrapped__, # pylint: disable=no-member + cls.take_snapshots = ttl_cache(ttl=interval)( # type: ignore[method-assign] + cls.take_snapshots.__wrapped__, # type: ignore[attr-defined] # pylint: disable=no-member ) @ttl_cache(ttl=1.0) - def take_snapshots(self): + def take_snapshots(self) -> list[Snapshot]: snapshots = [device.as_snapshot() for device in self.all_devices] for device in snapshots: @@ -155,10 +175,11 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes device.name = device.name.rpartition(' ')[-1] if device.bar1_memory_percent is not NA: device.bar1_memory_percent = round(device.bar1_memory_percent) - if device.bar1_memory_percent >= 100: - device.bar1_memory_percent_string = 'MAX' - else: - device.bar1_memory_percent_string = f'{round(device.bar1_memory_percent)}%' + device.bar1_memory_percent_string = ( + 'MAX' + if device.bar1_memory_percent >= 100 + else f'{round(device.bar1_memory_percent)}%' + ) else: device.name = cut_string(device.name, maxlen=18, padstr='..', align='right') device.current_driver_model = device.current_driver_model.replace('WDM', 'TCC') @@ -180,13 +201,13 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes return snapshots - def _snapshot_target(self): + def _snapshot_target(self) -> None: self._daemon_running.wait() while self._daemon_running.is_set(): self.take_snapshots() time.sleep(self.SNAPSHOT_INTERVAL) - def header_lines(self, compact=None): + def header_lines(self, compact: bool | None = None) -> list[str]: if compact is None: compact = self.compact @@ -218,7 +239,7 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes '│ Fan Temp Perf Pwr:Usage/Cap│ Memory-Usage │ GPU-Util Compute M. │', ), ) - if WINDOWS: + if IS_WINDOWS: header[-2] = header[-2].replace('Persistence-M', ' TCC/WDDM ') if self.support_mig: header[-2] = header[-2].replace('Volatile Uncorr. ECC', 'MIG M. Uncorr. ECC') @@ -235,7 +256,7 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes ) return header - def frame_lines(self, compact=None): + def frame_lines(self, compact: bool | None = None) -> list[str]: if compact is None: compact = self.compact @@ -273,7 +294,7 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes return frame - def poke(self): + def poke(self) -> None: if not self._daemon_running.is_set(): self._daemon_running.set() self._snapshot_daemon.start() @@ -283,7 +304,7 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes super().poke() # pylint: disable-next=too-many-locals,too-many-branches,too-many-statements - def draw(self): + def draw(self) -> None: self.color_reset() if self.need_redraw: @@ -300,7 +321,7 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes remaining_width = self.width - 79 draw_bars = self.width >= 100 try: - selected_device = self.parent.selection.process.device + selected_device = self.parent.selection.process.device # type: ignore[union-attr] except AttributeError: selected_device = None @@ -313,7 +334,7 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes ): y_start += 1 - attr = 0 + attr: int | str = 0 if selected_device is not None: if device.real == selected_device: attr = 'bold' @@ -403,16 +424,16 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes y_start += len(fmts) prev_device_index = device.tuple_index - def destroy(self): + def destroy(self) -> None: super().destroy() self._daemon_running.clear() - def print_width(self): + def print_width(self) -> int: if self.device_count > 0 and self.width >= 100: return self.width return 79 - def print(self): # pylint: disable=too-many-locals,too-many-branches + def print(self) -> None: # pylint: disable=too-many-locals,too-many-branches lines = [time.strftime('%a %b %d %H:%M:%S %Y'), *self.header_lines(compact=False)] if self.device_count > 0: @@ -426,8 +447,8 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes '├───────────────────────────────┼──────────────────────┼──────────────────────┤', ) - def colorize(s): - if len(s) > 0: + def colorize(s: str) -> str: + if s: # pylint: disable-next=cell-var-from-loop return colored(s, device.display_color) # noqa: B023 return '' @@ -494,15 +515,15 @@ class DevicePanel(Displayable): # pylint: disable=too-many-instance-attributes y_start += len(matrix) prev_device_index = device.tuple_index - lines = '\n'.join(lines) - if self.ascii: - lines = lines.translate(self.ASCII_TRANSTABLE) + content = '\n'.join(lines) + if self.no_unicode: + content = content.translate(self.ASCII_TRANSTABLE) try: - print(lines) + print(content) except UnicodeError: - print(lines.translate(self.ASCII_TRANSTABLE)) + print(content.translate(self.ASCII_TRANSTABLE)) - def press(self, key): + def press(self, key: int) -> bool: self.root.keymaps.use_keymap('device') - self.root.press(key) + return self.root.press(key) diff --git a/nvitop/tui/screens/main/host.py b/nvitop/tui/screens/main/panels/host.py similarity index 78% rename from nvitop/tui/screens/main/host.py rename to nvitop/tui/screens/main/panels/host.py index f0bc1cf..de61296 100644 --- a/nvitop/tui/screens/main/host.py +++ b/nvitop/tui/screens/main/panels/host.py @@ -3,48 +3,70 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import threading import time +from typing import TYPE_CHECKING, ClassVar from nvitop.tui.library import ( NA, BufferedHistoryGraph, Device, - Displayable, GiB, + HistoryGraph, + MigDevice, + NaType, bytes2human, colored, host, make_bar, timedelta2human, ) +from nvitop.tui.screens.main.panels.base import BasePanel -class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes - NAME = 'host' - SNAPSHOT_INTERVAL = 0.5 +if TYPE_CHECKING: + import curses - def __init__(self, devices, compact, win, root): + from nvitop.tui.tui import TUI + + +__all__ = ['HostPanel'] + + +class HostPanel(BasePanel): # pylint: disable=too-many-instance-attributes + NAME: ClassVar[str] = 'host' + SNAPSHOT_INTERVAL: ClassVar[float] = 0.5 + + def __init__( + self, + devices: list[Device | MigDevice], + compact: bool, + *, + win: curses.window | None, + root: TUI, + ) -> None: super().__init__(win, root) - self.devices = devices - self.device_count = len(self.devices) + self.devices: list[Device | MigDevice] = devices + self.device_count: int = len(self.devices) if win is not None: - self.average_gpu_memory_percent = None - self.average_gpu_utilization = None + self.average_gpu_memory_percent: HistoryGraph | None = None + self.average_gpu_utilization: HistoryGraph | None = None self.enable_history() - self._compact = compact - self.width = max(79, root.width) - self.full_height = 12 - self.compact_height = 2 - self.height = self.compact_height if compact else self.full_height + self._compact: bool = compact + self.width: int = max(79, root.width) + self.full_height: int = 12 + self.compact_height: int = 2 + self.height: int = self.compact_height if compact else self.full_height - self.cpu_percent = None - self.load_average = None - self.virtual_memory = None - self.swap_memory = None + self.cpu_percent: float = NA # type: ignore[assignment] + self.load_average: tuple[float, float, float] = (NA, NA, NA) # type: ignore[assignment] + self.virtual_memory: host.VirtualMemory = host.VirtualMemory() + self.swap_memory: host.SwapMemory = host.SwapMemory() self._snapshot_daemon = threading.Thread( name='host-snapshot-daemon', target=self._snapshot_target, @@ -53,37 +75,37 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes self._daemon_running = threading.Event() @property - def width(self): + def width(self) -> int: return self._width @width.setter - def width(self, value): + def width(self, value: int) -> None: width = max(79, value) if self._width != width: if self.visible: self.need_redraw = True graph_width = max(width - 80, 20) if self.win is not None: - self.average_gpu_memory_percent.width = graph_width - self.average_gpu_utilization.width = graph_width + self.average_gpu_memory_percent.width = graph_width # type: ignore[union-attr] + self.average_gpu_utilization.width = graph_width # type: ignore[union-attr] for device in self.devices: - device.memory_percent.history.width = graph_width - device.gpu_utilization.history.width = graph_width + device.memory_percent.history.width = graph_width # type: ignore[attr-defined] + device.gpu_utilization.history.width = graph_width # type: ignore[attr-defined] self._width = width @property - def compact(self): - return self._compact or self.ascii + def compact(self) -> bool: + return self._compact or self.no_unicode @compact.setter - def compact(self, value): - value = value or self.ascii + def compact(self, value: bool) -> None: + value = value or self.no_unicode if self._compact != value: self.need_redraw = True self._compact = value self.height = self.compact_height if self.compact else self.full_height - def enable_history(self): + def enable_history(self) -> None: host.cpu_percent = BufferedHistoryGraph( interval=1.0, width=77, @@ -115,11 +137,11 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes format='{:.1f}%'.format, )(host.swap_memory, get_value=lambda sm: sm.percent) - def percentage(x): + def percentage(x: float | NaType) -> str: return f'{x:.1f}%' if x is not NA else NA - def enable_history(device): - device.memory_percent = BufferedHistoryGraph( + def enable_history(device: Device) -> None: + device.memory_percent = BufferedHistoryGraph( # type: ignore[method-assign] interval=1.0, width=20, height=5, @@ -129,7 +151,7 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes dynamic_bound=False, format=lambda x: f'GPU {device.display_index} MEM: {percentage(x)}', )(device.memory_percent) - device.gpu_utilization = BufferedHistoryGraph( + device.gpu_utilization = BufferedHistoryGraph( # type: ignore[method-assign] interval=1.0, width=20, height=5, @@ -166,21 +188,21 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes ) @classmethod - def set_snapshot_interval(cls, interval): + def set_snapshot_interval(cls, interval: float) -> None: assert interval > 0.0 interval = float(interval) cls.SNAPSHOT_INTERVAL = min(interval / 3.0, 0.5) - def take_snapshots(self): + def take_snapshots(self) -> None: host.cpu_percent() host.virtual_memory() host.swap_memory() self.load_average = host.load_average() self.cpu_percent = host.cpu_percent.history.last_value - self.virtual_memory = host.virtual_memory.history.last_retval - self.swap_memory = host.swap_memory.history.last_retval + self.virtual_memory = host.virtual_memory.history.last_retval # type: ignore[attr-defined] + self.swap_memory = host.swap_memory.history.last_retval # type: ignore[attr-defined] total_memory_used = 0 total_memory_total = 0 @@ -195,20 +217,22 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes if gpu_utilization is not NA: gpu_utilizations.append(float(gpu_utilization)) if total_memory_total > 0: - self.average_gpu_memory_percent.add(100.0 * total_memory_used / total_memory_total) + avg = 100.0 * total_memory_used / total_memory_total + self.average_gpu_memory_percent.add(avg) # type: ignore[union-attr] if len(gpu_utilizations) > 0: - self.average_gpu_utilization.add(sum(gpu_utilizations) / len(gpu_utilizations)) + avg = sum(gpu_utilizations) / len(gpu_utilizations) + self.average_gpu_utilization.add(avg) # type: ignore[union-attr] - def _snapshot_target(self): + def _snapshot_target(self) -> None: self._daemon_running.wait() while self._daemon_running.is_set(): self.take_snapshots() time.sleep(self.SNAPSHOT_INTERVAL) - def frame_lines(self, compact=None): + def frame_lines(self, compact: bool | None = None) -> list[str]: if compact is None: compact = self.compact - if compact or self.ascii: + if compact or self.no_unicode: return [] remaining_width = self.width - 79 @@ -243,7 +267,7 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes return frame - def poke(self): + def poke(self) -> None: if not self._daemon_running.is_set(): self._daemon_running.set() self._snapshot_daemon.start() @@ -251,7 +275,7 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes super().poke() - def draw(self): # pylint: disable=too-many-locals,too-many-branches,too-many-statements + def draw(self) -> None: # pylint: disable=too-many-locals,too-many-branches,too-many-statements self.color_reset() load_average = 'Load Average: {} {} {}'.format( @@ -325,18 +349,18 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes self.addstr(y, self.x + 1, line) self.color(fg='magenta') - for y, line in enumerate(host.virtual_memory.history.graph, start=self.y + 6): + for y, line in enumerate(host.virtual_memory.history.graph, start=self.y + 6): # type: ignore[attr-defined] self.addstr(y, self.x + 1, line) self.color(fg='blue') - for y, line in enumerate(host.swap_memory.history.graph, start=self.y + 10): + for y, line in enumerate(host.swap_memory.history.graph, start=self.y + 10): # type: ignore[attr-defined] self.addstr(y, self.x + 1, line) if self.width >= 100: if self.device_count > 1 and self.parent.selection.is_set(): - device = self.parent.selection.process.device - gpu_memory_percent = device.memory_percent.history - gpu_utilization = device.gpu_utilization.history + device = self.parent.selection.process.device # type: ignore[union-attr] + gpu_memory_percent = device.memory_percent.history # type: ignore[union-attr] + gpu_utilization = device.gpu_utilization.history # type: ignore[union-attr] else: gpu_memory_percent = self.average_gpu_memory_percent gpu_utilization = self.average_gpu_utilization @@ -362,39 +386,41 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes self.addstr( self.y + 9, self.x + 1, - f' MEM: {bytes2human(self.virtual_memory.used, min_unit=GiB)} ({host.virtual_memory.history}) ', + ( + f' MEM: {bytes2human(self.virtual_memory.used, min_unit=GiB)} ' + f'({host.virtual_memory.history}) ' # type: ignore[attr-defined] + ), ) self.addstr( self.y + 10, self.x + 1, - f' SWP: {bytes2human(self.swap_memory.used, min_unit=GiB)} ({host.swap_memory.history}) ', + ( + f' SWP: {bytes2human(self.swap_memory.used, min_unit=GiB)} ' + f'({host.swap_memory.history}) ' # type: ignore[attr-defined] + ), ) if self.width >= 100: self.addstr(self.y, self.x + 79, f' {gpu_memory_percent} ') self.addstr(self.y + 10, self.x + 79, f' {gpu_utilization} ') - def destroy(self): + def destroy(self) -> None: super().destroy() self._daemon_running.clear() - def print_width(self): + def print_width(self) -> int: if self.device_count > 0 and self.width >= 100: return self.width return 79 - def print(self): + def print(self) -> None: self.cpu_percent = host.cpu_percent() self.virtual_memory = host.virtual_memory() self.swap_memory = host.swap_memory() self.load_average = host.load_average() - if self.load_average is not None: - load_average = tuple( - f'{value:5.2f}'[:5] if value < 10000.0 else '9999+' for value in self.load_average - ) - else: - load_average = (NA,) * 3 - load_average = 'Load Average: {} {} {}'.format(*load_average) + load_average = 'Load Average: {} {} {}'.format( + *(f'{value:5.2f}'[:5] if value < 10000.0 else '9999+' for value in self.load_average), + ) width_right = len(load_average) + 4 width_left = self.width - 2 - width_right @@ -428,7 +454,7 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes ] lines = '\n'.join(lines) - if self.ascii: + if self.no_unicode: lines = lines.translate(self.ASCII_TRANSTABLE) try: @@ -436,6 +462,6 @@ class HostPanel(Displayable): # pylint: disable=too-many-instance-attributes except UnicodeError: print(lines.translate(self.ASCII_TRANSTABLE)) - def press(self, key): + def press(self, key: int) -> bool: self.root.keymaps.use_keymap('host') - self.root.press(key) + return self.root.press(key) diff --git a/nvitop/tui/screens/main/process.py b/nvitop/tui/screens/main/panels/process.py similarity index 69% rename from nvitop/tui/screens/main/process.py rename to nvitop/tui/screens/main/panels/process.py index 5535b3b..6302d39 100644 --- a/nvitop/tui/screens/main/process.py +++ b/nvitop/tui/screens/main/panels/process.py @@ -9,30 +9,54 @@ import itertools import threading import time from operator import attrgetter, xor -from typing import TYPE_CHECKING, Any, NamedTuple +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple from nvitop.tui.library import ( HOSTNAME, + IS_SUPERUSER, + IS_WINDOWS, + IS_WSL, LARGE_INTEGER, - SUPERUSER, USER_CONTEXT, USERNAME, - WINDOWS, - WSL, + Device, Displayable, GpuProcess, + MigDevice, MouseEvent, Selection, + Snapshot, WideString, colored, cut_string, ttl_cache, wcslen, ) +from nvitop.tui.screens.main.panels.base import BaseSelectablePanel if TYPE_CHECKING: - from collections.abc import Callable + import curses + from collections.abc import Callable, Iterable, Mapping + + from nvitop.tui.tui import TUI + + +__all__ = ['OrderName', 'ProcessPanel'] + + +OrderName = Literal[ + 'natural', + 'pid', + 'username', + 'gpu_memory', + 'sm_utilization', + 'gpu_memory_utilization', + 'cpu_percent', + 'memory_percent', + 'time', +] class Order(NamedTuple): @@ -40,144 +64,185 @@ class Order(NamedTuple): reverse: bool offset: int column: str - previous: str - next: str + previous: OrderName + next: OrderName bind_key: str -class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes - NAME = 'process' - SNAPSHOT_INTERVAL = 0.5 +class ProcessPanel(BaseSelectablePanel): # pylint: disable=too-many-instance-attributes + NAME: ClassVar[str] = 'process' + SNAPSHOT_INTERVAL: ClassVar[float] = 0.5 - ORDERS = { - 'natural': Order( - key=attrgetter('device.tuple_index', '_gone', 'username', 'pid'), - reverse=False, - offset=3, - column='ID', - previous='time', - next='pid', - bind_key='n', - ), - 'pid': Order( - key=attrgetter('_gone', 'pid', 'device.tuple_index'), - reverse=False, - offset=10, - column='PID', - previous='natural', - next='username', - bind_key='p', - ), - 'username': Order( - key=attrgetter('_gone', 'username', 'pid', 'device.tuple_index'), - reverse=False, - offset=19, - column='USER', - previous='pid', - next='gpu_memory', - bind_key='u', - ), - 'gpu_memory': Order( - key=attrgetter( - '_gone', - 'gpu_memory', - 'gpu_sm_utilization', - 'cpu_percent', - 'pid', - 'device.tuple_index', + ORDERS: ClassVar[Mapping[OrderName, Order]] = MappingProxyType( + { + 'natural': Order( + key=attrgetter( + 'device.tuple_index', + '_gone', + 'username', + 'pid', + ), + reverse=False, + offset=3, + column='ID', + previous='time', + next='pid', + bind_key='n', ), - reverse=True, - offset=25, - column='GPU-MEM', - previous='username', - next='sm_utilization', - bind_key='g', - ), - 'sm_utilization': Order( - key=attrgetter( - '_gone', - 'gpu_sm_utilization', - 'gpu_memory', - 'cpu_percent', - 'pid', - 'device.tuple_index', + 'pid': Order( + key=attrgetter( + '_gone', + 'pid', + 'device.tuple_index', + ), + reverse=False, + offset=10, + column='PID', + previous='natural', + next='username', + bind_key='p', ), - reverse=True, - offset=34, - column='SM', - previous='gpu_memory', - next='gpu_memory_utilization', - bind_key='s', - ), - 'gpu_memory_utilization': Order( - key=attrgetter( - '_gone', - 'gpu_memory_utilization', - 'gpu_memory', - 'cpu_percent', - 'pid', - 'device.tuple_index', + 'username': Order( + key=attrgetter( + '_gone', + 'username', + 'pid', + 'device.tuple_index', + ), + reverse=False, + offset=19, + column='USER', + previous='pid', + next='gpu_memory', + bind_key='u', ), - reverse=True, - offset=38, - column='GMBW', - previous='gpu_memory', - next='cpu_percent', - bind_key='b', - ), - 'cpu_percent': Order( - key=attrgetter('_gone', 'cpu_percent', 'memory_percent', 'pid', 'device.tuple_index'), - reverse=True, - offset=44, - column='%CPU', - previous='gpu_memory_utilization', - next='memory_percent', - bind_key='c', - ), - 'memory_percent': Order( - key=attrgetter('_gone', 'memory_percent', 'cpu_percent', 'pid', 'device.tuple_index'), - reverse=True, - offset=50, - column='%MEM', - previous='cpu_percent', - next='time', - bind_key='m', - ), - 'time': Order( - key=attrgetter('_gone', 'running_time', 'pid', 'device.tuple_index'), - reverse=True, - offset=56, - column='TIME', - previous='memory_percent', - next='natural', - bind_key='t', - ), - } + 'gpu_memory': Order( + key=attrgetter( + '_gone', + 'gpu_memory', + 'gpu_sm_utilization', + 'cpu_percent', + 'pid', + 'device.tuple_index', + ), + reverse=True, + offset=25, + column='GPU-MEM', + previous='username', + next='sm_utilization', + bind_key='g', + ), + 'sm_utilization': Order( + key=attrgetter( + '_gone', + 'gpu_sm_utilization', + 'gpu_memory', + 'cpu_percent', + 'pid', + 'device.tuple_index', + ), + reverse=True, + offset=34, + column='SM', + previous='gpu_memory', + next='gpu_memory_utilization', + bind_key='s', + ), + 'gpu_memory_utilization': Order( + key=attrgetter( + '_gone', + 'gpu_memory_utilization', + 'gpu_memory', + 'cpu_percent', + 'pid', + 'device.tuple_index', + ), + reverse=True, + offset=38, + column='GMBW', + previous='sm_utilization', + next='cpu_percent', + bind_key='b', + ), + 'cpu_percent': Order( + key=attrgetter( + '_gone', + 'cpu_percent', + 'memory_percent', + 'pid', + 'device.tuple_index', + ), + reverse=True, + offset=44, + column='%CPU', + previous='gpu_memory_utilization', + next='memory_percent', + bind_key='c', + ), + 'memory_percent': Order( + key=attrgetter( + '_gone', + 'memory_percent', + 'cpu_percent', + 'pid', + 'device.tuple_index', + ), + reverse=True, + offset=50, + column='%MEM', + previous='cpu_percent', + next='time', + bind_key='m', + ), + 'time': Order( + key=attrgetter( + '_gone', + 'running_time', + 'pid', + 'device.tuple_index', + ), + reverse=True, + offset=56, + column='TIME', + previous='memory_percent', + next='natural', + bind_key='t', + ), + }, + ) # pylint: disable-next=too-many-arguments - def __init__(self, devices, compact, filters, *, win, root): + def __init__( + self, + devices: list[Device | MigDevice], + compact: bool, + filters: Iterable[Callable[[Snapshot], bool]], + *, + win: curses.window | None, + root: TUI, + ) -> None: super().__init__(win, root) - self.devices = devices + self.devices: list[Device | MigDevice] = devices - self._compact = compact - self.width = max(79, root.width) + self._compact: bool = compact + self.width: int = max(79, root.width) self.height = self._full_height = self.compact_height = 7 - self.filters = [None, *filters] + self.filters: list[Callable[[Snapshot], bool] | None] = [None, *filters] - self.host_headers = ['%CPU', '%MEM', 'TIME', 'COMMAND'] + self.host_headers: list[str] = ['%CPU', '%MEM', 'TIME', 'COMMAND'] - self.selection = Selection(panel=self) - self.host_offset = -1 - self.y_mouse = None + self.selection: Selection = Selection(self) + self.host_offset: int = -1 + self.y_mouse: int | None = None - self._order = 'natural' - self.reverse = False + self._order: OrderName = 'natural' + self.reverse: bool = False - self.has_snapshots = False - self._snapshot_buffer = None - self._snapshots = [] + self.has_snapshots: int = False + self._snapshot_buffer: list[Snapshot] | None = None + self._snapshots: list[Snapshot] = [] self.snapshot_lock = threading.Lock() self._snapshot_daemon = threading.Thread( name='process-snapshot-daemon', @@ -187,22 +252,22 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes self._daemon_running = threading.Event() @property - def width(self): + def width(self) -> int: return self._width @width.setter - def width(self, value): + def width(self, value: int) -> None: width = max(79, value) if self._width != width and self.visible: self.need_redraw = True self._width = width @property - def compact(self): + def compact(self) -> bool: return self._compact or self.order != 'natural' @compact.setter - def compact(self, value): + def compact(self, value: bool) -> None: if self._compact != value: self.need_redraw = True self._compact = value @@ -216,29 +281,29 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes self.height = self.compact_height if self.compact else self.full_height @property - def full_height(self): + def full_height(self) -> int: return self._full_height if self.order == 'natural' else self.compact_height @full_height.setter - def full_height(self, value): + def full_height(self, value: int) -> None: self._full_height = value @property - def order(self): + def order(self) -> OrderName: return self._order @order.setter - def order(self, value): + def order(self, value: OrderName) -> None: if self._order != value: self._order = value self.height = self.compact_height if self.compact else self.full_height @property - def snapshots(self): + def snapshots(self) -> list[Snapshot]: return self._snapshots @snapshots.setter - def snapshots(self, snapshots): + def snapshots(self, snapshots: list[Snapshot] | None) -> None: if snapshots is None: return self.has_snapshots = True @@ -273,28 +338,28 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes for i, process in enumerate(snapshots): if process._ident == identity: # pylint: disable=protected-access self.selection.index = i - self.selection.process = process + self.selection.process = process # type: ignore[assignment] break @classmethod - def set_snapshot_interval(cls, interval): + def set_snapshot_interval(cls, interval: float) -> None: assert interval > 0.0 interval = float(interval) cls.SNAPSHOT_INTERVAL = min(interval / 3.0, 1.0) - cls.take_snapshots = ttl_cache(ttl=interval)( - cls.take_snapshots.__wrapped__, # pylint: disable=no-member + cls.take_snapshots = ttl_cache(ttl=interval)( # type: ignore[method-assign] + cls.take_snapshots.__wrapped__, # type: ignore[attr-defined] # pylint: disable=no-member ) - def ensure_snapshots(self): + def ensure_snapshots(self) -> None: if not self.has_snapshots: self.snapshots = self.take_snapshots() @ttl_cache(ttl=2.0) - def take_snapshots(self): + def take_snapshots(self) -> list[Snapshot]: snapshots = GpuProcess.take_snapshots(self.processes, failsafe=True) for condition in self.filters: - snapshots = filter(condition, snapshots) + snapshots = filter(condition, snapshots) # type: ignore[assignment] snapshots = list(snapshots) time_length = max(4, max((len(p.running_time_human) for p in snapshots), default=4)) @@ -314,13 +379,13 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes return snapshots - def _snapshot_target(self): + def _snapshot_target(self) -> None: self._daemon_running.wait() while self._daemon_running.is_set(): self.take_snapshots() time.sleep(self.SNAPSHOT_INTERVAL) - def header_lines(self): + def header_lines(self) -> list[str]: header = [ '╒' + '═' * (self.width - 2) + '╕', '│ {} │'.format('Processes:'.ljust(self.width - 4)), @@ -331,7 +396,7 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes ] if len(self.snapshots) == 0: if self.has_snapshots: - message = ' No running processes found{} '.format(' (in WSL)' if WSL else '') + message = ' No running processes found{} '.format(' (in WSL)' if IS_WSL else '') else: message = ' Gathering process status...' header.extend( @@ -340,22 +405,22 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes return header @property - def processes(self): + def processes(self) -> list[GpuProcess]: return list( - itertools.chain.from_iterable(device.processes().values() for device in self.devices), + itertools.chain.from_iterable(device.processes().values() for device in self.devices), # type: ignore[misc] ) - def poke(self): + def poke(self) -> None: if not self._daemon_running.is_set(): self._daemon_running.set() self._snapshot_daemon.start() - self.snapshots = self._snapshot_buffer + self.snapshots = self._snapshot_buffer # type: ignore[assignment] self.selection.within_window = False if len(self.snapshots) > 0 and self.selection.is_set(): y = self.y + 5 - prev_device_index = None + prev_device_index: int | None = None for process in self.snapshots: device_index = process.device.physical_index if prev_device_index != device_index: @@ -365,13 +430,14 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes if self.selection.is_same(process): self.selection.within_window = ( - self.root.y <= y < self.root.termsize[0] and self.width >= 79 + self.root.y <= y < self.root.termsize[0] # type: ignore[index] + and self.width >= 79 ) if not self.selection.within_window: if y < self.root.y: self.parent.y += self.root.y - y - elif y >= self.root.termsize[0]: - self.parent.y -= y - self.root.termsize[0] + 1 + elif y >= self.root.termsize[0]: # type: ignore[index] + self.parent.y -= y - self.root.termsize[0] + 1 # type: ignore[index] self.parent.update_size(self.root.termsize) self.need_redraw = True break @@ -379,11 +445,11 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes super().poke() - def draw(self): # pylint: disable=too-many-locals,too-many-branches,too-many-statements + def draw(self) -> None: # pylint: disable=too-many-locals,too-many-branches,too-many-statements self.color_reset() if self.need_redraw: - if SUPERUSER: + if IS_SUPERUSER: self.addstr(self.y, self.x + 1, '!CAUTION: SUPERUSER LOGGED-IN.') self.color_at(self.y, self.x + 1, width=1, fg='red', attr='blink') self.color_at(self.y, self.x + 2, width=29, fg='yellow', attr='italic') @@ -392,7 +458,7 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes self.addstr(y, self.x, line) context_width = wcslen(USER_CONTEXT) - if not WINDOWS or len(USER_CONTEXT) == context_width: + if not IS_WINDOWS or len(USER_CONTEXT) == context_width: # Do not support windows-curses with wide characters username_width = wcslen(USERNAME) hostname_width = wcslen(HOSTNAME) @@ -403,7 +469,7 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes self.y + 2, self.x + offset, width=username_width, - fg=('yellow' if SUPERUSER else 'magenta'), + fg=('yellow' if IS_SUPERUSER else 'magenta'), attr='bold', ) self.color_at( @@ -474,7 +540,7 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes self.selection.within_window = False if len(self.snapshots) > 0: y = self.y + 5 - prev_device_index = None + prev_device_index: int | None = None prev_device_display_index = None color = -1 for process in self.snapshots: @@ -522,7 +588,7 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes self.addstr(y, self.x + 44, process.command) if y == self.y_mouse: - self.selection.process = process + self.selection.process = process # type: ignore[assignment] hint = True if self.selection.is_same(process): @@ -534,10 +600,11 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes attr='bold | reverse', ) self.selection.within_window = ( - self.root.y <= y < self.root.termsize[0] and self.width >= 79 + self.root.y <= y < self.root.termsize[0] # type: ignore[index] + and self.width >= 79 ) else: - owned = str(process.username) == USERNAME or SUPERUSER + owned = IS_SUPERUSER or str(process.username) == USERNAME if self.selection.is_same_on_host(process): self.addstr(y, self.x + 1, '=', self.get_fg_bg_attr(attr='bold | blink')) self.color_at(y, self.x + 2, width=3, fg=color) @@ -562,7 +629,7 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes self.selection.clear() elif self.has_snapshots: - message = ' No running processes found{} '.format(' (in WSL)' if WSL else '') + message = ' No running processes found{} '.format(' (in WSL)' if IS_WSL else '') self.addstr(self.y + 5, self.x, f'│ {message.ljust(self.width - 4)} │') text_offset = self.x + self.width - 47 @@ -579,29 +646,29 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes else: self.addstr(self.y, text_offset, ' ' * 47) - def finalize(self): + def finalize(self) -> None: self.y_mouse = None super().finalize() - def destroy(self): + def destroy(self) -> None: super().destroy() self._daemon_running.clear() - def print_width(self): + def print_width(self) -> int: self.ensure_snapshots() return min( self.width, max((45 + len(process.host_info) for process in self.snapshots), default=79), ) - def print(self): + def print(self) -> None: self.ensure_snapshots() lines = ['', *self.header_lines()] lines[2] = ''.join( ( lines[2][: -2 - wcslen(USER_CONTEXT)], - colored(USERNAME, color=('yellow' if SUPERUSER else 'magenta'), attrs=('bold',)), + colored(USERNAME, color=('yellow' if IS_SUPERUSER else 'magenta'), attrs=('bold',)), colored('@', attrs=('bold',)), colored(HOSTNAME, color='green', attrs=('bold',)), lines[2][-2:], @@ -611,7 +678,7 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes if len(self.snapshots) > 0: key, reverse, *_ = self.ORDERS['natural'] self.snapshots.sort(key=key, reverse=reverse) - prev_device_index = None + prev_device_index: int | None = None prev_device_display_index = None color = None for process in self.snapshots: @@ -637,21 +704,21 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes WideString(host_info).ljust(self.width - 45)[: self.width - 45], ) if process.is_zombie or process.no_permissions or process.is_gone: - info = info.split(process.command) - if process.username != USERNAME and not SUPERUSER: - info = (colored(item, attrs=('dark',)) for item in info) + info_segments = info.split(process.command) + if not IS_SUPERUSER and process.username != USERNAME: + info_segments = [colored(item, attrs=('dark',)) for item in info] info = colored( process.command, color=('red' if process.is_gone else 'yellow'), - ).join(info) - elif process.username != USERNAME and not SUPERUSER: + ).join(info_segments) + elif not IS_SUPERUSER and process.username != USERNAME: info = colored(info, attrs=('dark',)) lines.append('│{} {} │'.format(colored(f'{device_display_index:>4}', color), info)) lines.append('╘' + '═' * (self.width - 2) + '╛') lines = '\n'.join(lines) - if self.ascii: + if self.no_unicode: lines = lines.translate(self.ASCII_TRANSTABLE) try: @@ -659,11 +726,11 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes except UnicodeError: print(lines.translate(self.ASCII_TRANSTABLE)) - def press(self, key): + def press(self, key: int) -> bool: self.root.keymaps.use_keymap('process') - self.root.press(key) + return self.root.press(key) - def click(self, event): + def click(self, event: MouseEvent) -> bool: if event.pressed(1) or event.pressed(3) or event.clicked(1) or event.clicked(3): self.y_mouse = event.y return True @@ -675,7 +742,7 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes self.selection.move(direction=direction) return True - def __contains__(self, item): + def __contains__(self, item: Displayable | MouseEvent | tuple[int, int]) -> bool: if self.parent.visible and isinstance(item, MouseEvent): return True return super().__contains__(item) diff --git a/nvitop/tui/screens/metrics.py b/nvitop/tui/screens/metrics.py index f2de09b..c58b3c0 100644 --- a/nvitop/tui/screens/metrics.py +++ b/nvitop/tui/screens/metrics.py @@ -4,20 +4,24 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring # pylint: disable=invalid-name +from __future__ import annotations + import itertools import threading import time from collections import OrderedDict +from typing import TYPE_CHECKING, ClassVar from nvitop.tui.library import ( HOSTNAME, + IS_SUPERUSER, + IS_WINDOWS, NA, - SUPERUSER, USER_CONTEXT, USERNAME, BufferedHistoryGraph, - Displayable, GpuProcess, + HistoryGraph, Selection, WideString, bytes2human, @@ -25,17 +29,28 @@ from nvitop.tui.library import ( host, wcslen, ) +from nvitop.tui.screens.base import BaseSelectableScreen -def get_yticks(history, y_offset): # pylint: disable=too-many-branches,too-many-locals +if TYPE_CHECKING: + import curses + + from nvitop.tui.tui import TUI + + +__all__ = ['ProcessMetricsScreen'] + + +# pylint: disable-next=too-many-branches,too-many-locals +def get_yticks(history: HistoryGraph, y_offset: int) -> list[tuple[int, int]]: height = history.height baseline = history.baseline bound = history.bound max_bound = history.max_bound - scale = history.scale + scale: float = history.scale # type: ignore[attr-defined] upsidedown = history.upsidedown - def p2h_f(p): + def p2h_f(p: int) -> float: return 0.01 * scale * p * (max_bound - baseline) * (height - 1) / (bound - baseline) max_height = height - 2 @@ -77,20 +92,20 @@ def get_yticks(history, y_offset): # pylint: disable=too-many-branches,too-many return [(h + y_offset, p) for h, p in ticks] -class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-attributes - NAME = 'process-metrics' - SNAPSHOT_INTERVAL = 0.5 +class ProcessMetricsScreen(BaseSelectableScreen): # pylint: disable=too-many-instance-attributes + NAME: ClassVar[str] = 'process-metrics' + SNAPSHOT_INTERVAL: ClassVar[float] = 0.5 - def __init__(self, win, root): + def __init__(self, *, win: curses.window, root: TUI) -> None: super().__init__(win, root) - self.selection = Selection(panel=self) - self.used_gpu_memory = None - self.gpu_sm_utilization = None - self.cpu_percent = None - self.used_host_memory = None + self.selection: Selection = Selection(self) + self.used_gpu_memory: HistoryGraph | None = None + self.gpu_sm_utilization: HistoryGraph | None = None + self.cpu_percent: HistoryGraph | None = None + self.used_host_memory: HistoryGraph | None = None - self.enabled = False + self.enabled: bool = False self.snapshot_lock = threading.Lock() self._snapshot_daemon = threading.Thread( name='process-metrics-snapshot-daemon', @@ -101,17 +116,17 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at self.x, self.y = root.x, root.y self.width, self.height = root.width, root.height - self.left_width = max(20, (self.width - 3) // 2) - self.right_width = max(20, (self.width - 2) // 2) - self.upper_height = max(5, (self.height - 5 - 3) // 2) - self.lower_height = max(5, (self.height - 5 - 2) // 2) + self.left_width: int = max(20, (self.width - 3) // 2) + self.right_width: int = max(20, (self.width - 2) // 2) + self.upper_height: int = max(5, (self.height - 5 - 3) // 2) + self.lower_height: int = max(5, (self.height - 5 - 2) // 2) @property - def visible(self): + def visible(self) -> bool: return self._visible @visible.setter - def visible(self, value): + def visible(self, value: bool) -> None: if self._visible != value: self.need_redraw = True self._visible = value @@ -125,7 +140,7 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at else: self.focused = False - def enable(self, state=True): + def enable(self, state: bool = True) -> None: if not self.selection.is_set() or not state: self.disable() return @@ -135,57 +150,57 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at total_gpu_memory = self.process.device.memory_total() total_gpu_memory_human = bytes2human(total_gpu_memory) - def format_cpu_percent(value): - if value is NA: + def format_cpu_percent(value: float) -> str: + if value is NA: # type: ignore[comparison-overlap] return f'CPU: {value}' return f'CPU: {value:.1f}%' - def format_max_cpu_percent(value): - if value is NA: + def format_max_cpu_percent(value: float) -> str: + if value is NA: # type: ignore[comparison-overlap] return f'MAX CPU: {value}' return f'MAX CPU: {value:.1f}%' - def format_host_memory(value): - if value is NA: + def format_host_memory(value: float) -> str: + if value is NA: # type: ignore[comparison-overlap] return f'HOST-MEM: {value}' - return 'HOST-MEM: {} ({:.1f}%)'.format( # noqa: UP032 - bytes2human(value), - round(100.0 * value / total_host_memory, 1), + return ( + f'HOST-MEM: {bytes2human(value)} ' + f'({round(100.0 * value / total_host_memory, 1):.1f}%)' ) - def format_max_host_memory(value): - if value is NA: + def format_max_host_memory(value: float) -> str: + if value is NA: # type: ignore[comparison-overlap] return f'MAX HOST-MEM: {value}' - return 'MAX HOST-MEM: {} ({:.1f}%) / {}'.format( # noqa: UP032 - bytes2human(value), - round(100.0 * value / total_host_memory, 1), - total_host_memory_human, + return ( + f'MAX HOST-MEM: {bytes2human(value)} ' + f'({round(100.0 * value / total_host_memory, 1):.1f}%) ' + f'/ {total_host_memory_human}' ) - def format_gpu_memory(value): - if value is not NA and total_gpu_memory is not NA: - return 'GPU-MEM: {} ({:.1f}%)'.format( # noqa: UP032 - bytes2human(value), - round(100.0 * value / total_gpu_memory, 1), + def format_gpu_memory(value: float) -> str: + if value is not NA and total_gpu_memory is not NA: # type: ignore[comparison-overlap] + return ( + f'GPU-MEM: {bytes2human(value)} ' + f'({round(100.0 * value / total_gpu_memory, 1):.1f}%)' ) return f'GPU-MEM: {value}' - def format_max_gpu_memory(value): - if value is not NA and total_gpu_memory is not NA: - return 'MAX GPU-MEM: {} ({:.1f}%) / {}'.format( # noqa: UP032 - bytes2human(value), - round(100.0 * value / total_gpu_memory, 1), - total_gpu_memory_human, + def format_max_gpu_memory(value: float) -> str: + if value is not NA and total_gpu_memory is not NA: # type: ignore[comparison-overlap] + return ( + f'MAX GPU-MEM: {bytes2human(value)} ' + f'({round(100.0 * value / total_gpu_memory, 1):.1f}%) ' + f'/ {total_gpu_memory_human}' ) return f'MAX GPU-MEM: {value}' - def format_sm(value): - if value is NA: + def format_sm(value: float) -> str: + if value is NA: # type: ignore[comparison-overlap] return f'GPU-SM: {value}' return f'GPU-SM: {value:.1f}%' - def format_max_sm(value): - if value is NA: + def format_max_sm(value: float) -> str: + if value is NA: # type: ignore[comparison-overlap] return f'MAX GPU-SM: {value}' return f'MAX GPU-SM: {value:.1f}%' @@ -216,7 +231,7 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at ) self.used_gpu_memory = BufferedHistoryGraph( interval=1.0, - upperbound=total_gpu_memory or 1, + upperbound=total_gpu_memory or 1.0, # type: ignore[arg-type] width=self.right_width, height=self.upper_height, baseline=0.0, @@ -236,10 +251,10 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at format=format_sm, max_format=format_max_sm, ) - self.cpu_percent.scale = 0.1 - self.used_host_memory.scale = 1.0 - self.used_gpu_memory.scale = 1.0 - self.gpu_sm_utilization.scale = 1.0 + self.cpu_percent.scale = 0.1 # type: ignore[attr-defined] + self.used_host_memory.scale = 1.0 # type: ignore[attr-defined] + self.used_gpu_memory.scale = 1.0 # type: ignore[attr-defined] + self.gpu_sm_utilization.scale = 1.0 # type: ignore[attr-defined] self._daemon_running.set() try: @@ -251,7 +266,7 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at self.take_snapshots() self.update_size() - def disable(self): + def disable(self) -> None: with self.snapshot_lock: self._daemon_running.clear() self.enabled = False @@ -261,22 +276,22 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at self.gpu_sm_utilization = None @property - def process(self): - return self.selection.process + def process(self) -> GpuProcess: + return self.selection.process # type: ignore[return-value] @process.setter - def process(self, value): + def process(self, value: GpuProcess) -> None: self.selection.process = value self.enable() @classmethod - def set_snapshot_interval(cls, interval): + def set_snapshot_interval(cls, interval: float) -> None: assert interval > 0.0 interval = float(interval) cls.SNAPSHOT_INTERVAL = min(interval / 3.0, 1.0) - def take_snapshots(self): + def take_snapshots(self) -> None: with self.snapshot_lock: if not self.selection.is_set() or not self.enabled: return @@ -286,18 +301,22 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at self.process.update_gpu_status() snapshot = self.process.as_snapshot() + assert self.cpu_percent is not None + assert self.used_host_memory is not None + assert self.used_gpu_memory is not None + assert self.gpu_sm_utilization is not None self.cpu_percent.add(snapshot.cpu_percent) self.used_host_memory.add(snapshot.host_memory) self.used_gpu_memory.add(snapshot.gpu_memory) self.gpu_sm_utilization.add(snapshot.gpu_sm_utilization) - def _snapshot_target(self): + def _snapshot_target(self) -> None: while True: self._daemon_running.wait() self.take_snapshots() time.sleep(self.SNAPSHOT_INTERVAL) - def update_size(self, termsize=None): + def update_size(self, termsize: tuple[int, int] | None = None) -> tuple[int, int]: n_term_lines, n_term_cols = termsize = super().update_size(termsize=termsize) self.width = n_term_cols - self.x @@ -310,6 +329,10 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at with self.snapshot_lock: if self.enabled: + assert self.cpu_percent is not None + assert self.used_host_memory is not None + assert self.used_gpu_memory is not None + assert self.gpu_sm_utilization is not None self.cpu_percent.graph_size = (self.left_width, self.upper_height) self.used_host_memory.graph_size = (self.left_width, self.lower_height) self.used_gpu_memory.graph_size = (self.right_width, self.upper_height) @@ -317,7 +340,7 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at return termsize - def frame_lines(self): + def frame_lines(self) -> list[str]: line = '│' + ' ' * self.left_width + '│' + ' ' * self.right_width + '│' return [ '╒' + '═' * (self.width - 2) + '╕', @@ -332,7 +355,7 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at '╘' + '═' * self.left_width + '╧' + '═' * self.right_width + '╛', ] - def poke(self): + def poke(self) -> None: if self.visible and not self._daemon_running.is_set(): self._daemon_running.set() try: @@ -343,15 +366,20 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at super().poke() - def draw(self): # pylint: disable=too-many-statements,too-many-locals,too-many-branches + def draw(self) -> None: # pylint: disable=too-many-statements,too-many-locals,too-many-branches self.color_reset() + assert self.used_gpu_memory is not None + assert self.gpu_sm_utilization is not None + assert self.cpu_percent is not None + assert self.used_host_memory is not None + if self.need_redraw: for y, line in enumerate(self.frame_lines(), start=self.y): self.addstr(y, self.x, line) context_width = wcslen(USER_CONTEXT) - if not host.WINDOWS or len(USER_CONTEXT) == context_width: + if not IS_WINDOWS or len(USER_CONTEXT) == context_width: # Do not support windows-curses with wide characters username_width = wcslen(USERNAME) hostname_width = wcslen(HOSTNAME) @@ -362,7 +390,7 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at self.y + 1, self.x + offset, width=username_width, - fg=('yellow' if SUPERUSER else 'magenta'), + fg=('yellow' if IS_SUPERUSER else 'magenta'), attr='bold', ) self.color_at( @@ -589,10 +617,10 @@ class ProcessMetricsScreen(Displayable): # pylint: disable=too-many-instance-at self.addstr(y, x, f'├╴{p}% ') self.color_at(y, x, width=2, attr=0) - def destroy(self): + def destroy(self) -> None: super().destroy() self._daemon_running.clear() - def press(self, key): + def press(self, key: int) -> bool: self.root.keymaps.use_keymap('process-metrics') - self.root.press(key) + return self.root.press(key) diff --git a/nvitop/tui/screens/treeview.py b/nvitop/tui/screens/treeview.py index a7dab4f..ae2f457 100644 --- a/nvitop/tui/screens/treeview.py +++ b/nvitop/tui/screens/treeview.py @@ -3,41 +3,63 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import threading import time from collections import deque from functools import partial from itertools import islice +from typing import TYPE_CHECKING, Any, ClassVar from nvitop.tui.library import ( + IS_SUPERUSER, + IS_WSL, NA, - SUPERUSER, USERNAME, - Displayable, + Device, + GpuProcess, HostProcess, + MessageBox, + MouseEvent, Selection, Snapshot, WideString, host, - send_signal, ttl_cache, ) +from nvitop.tui.screens.base import BaseSelectableScreen + + +if TYPE_CHECKING: + import curses + from collections.abc import Iterable + from typing_extensions import Self # Python 3.11+ + + from nvitop.tui.tui import TUI + + +__all__ = ['TreeViewScreen'] class TreeNode: # pylint: disable=too-many-instance-attributes - def __init__(self, process, children=()): - self.process = process - self.parent = None - self.children = [] - self.devices = set() - self.children_set = set() - self.is_root = True - self.is_last = False - self.prefix = '' + def __init__( + self, + process: GpuProcess | HostProcess, + children: Iterable[Self] = (), + ) -> None: + self.process: GpuProcess | HostProcess = process + self.parent: TreeNode | None = None + self.children: list[Self] = [] + self.devices: set[Device] = set() + self.children_set: set[Self] = set() + self.is_root: bool = True + self.is_last: bool = False + self.prefix: str = '' for child in children: self.add(child) - def add(self, child): + def add(self, child: Self) -> None: if child in self.children_set: return self.children.append(child) @@ -45,19 +67,20 @@ class TreeNode: # pylint: disable=too-many-instance-attributes child.parent = self child.is_root = False - def __getattr__(self, name): + def __getattr__(self, name: str) -> Any: try: - return super().__getattr__(name) + return super().__getattr__(name) # type: ignore[misc] except AttributeError: return getattr(self.process, name) - def __eq__(self, other): - return self.process._ident == other.process._ident # pylint: disable=protected-access + def __eq__(self, other: object) -> bool: + # pylint: disable-next=protected-access + return self.process._ident == other.process._ident # type: ignore[attr-defined] - def __hash__(self): + def __hash__(self) -> int: return hash(self.process) - def as_snapshot(self): # pylint: disable=too-many-branches,too-many-statements + def as_snapshot(self) -> None: # pylint: disable=too-many-branches,too-many-statements if not isinstance(self.process, Snapshot): with self.process.oneshot(): try: @@ -73,6 +96,7 @@ class TreeNode: # pylint: disable=too-many-instance-attributes except host.PsutilError: command = 'No Such Process' + cpu_percent_string: str try: cpu_percent = self.process.cpu_percent() except host.PsutilError: @@ -87,15 +111,15 @@ class TreeNode: # pylint: disable=too-many-instance-attributes else: cpu_percent_string = '9999+%' + memory_percent_string: str try: memory_percent = self.process.memory_percent() except host.PsutilError: memory_percent = memory_percent_string = NA else: - if memory_percent is not NA: - memory_percent_string = f'{memory_percent:.1f}%' - else: - memory_percent_string = NA + memory_percent_string = ( + f'{memory_percent:.1f}%' if memory_percent is not NA else NA + ) try: num_threads = self.process.num_threads() @@ -107,7 +131,7 @@ class TreeNode: # pylint: disable=too-many-instance-attributes except host.PsutilError: running_time_human = NA - self.process = Snapshot( + self.process = Snapshot( # type: ignore[assignment] real=self.process, pid=self.process.pid, username=username, @@ -134,7 +158,7 @@ class TreeNode: # pylint: disable=too-many-instance-attributes child.is_last = False self.children[-1].is_last = True - def set_prefix(self, prefix=''): + def set_prefix(self, prefix: str = '') -> None: if self.is_root: self.prefix = '' else: @@ -144,16 +168,18 @@ class TreeNode: # pylint: disable=too-many-instance-attributes child.set_prefix(prefix) @classmethod - def merge(cls, leaves): # pylint: disable=too-many-branches - nodes = {} + def merge( # pylint: disable=too-many-branches + cls, + leaves: list[Snapshot | GpuProcess | HostProcess], + ) -> list[Self]: + nodes: dict[int, Self] = {} for process in leaves: - if isinstance(process, Snapshot): - process = process.real + real_process = process.real if isinstance(process, Snapshot) else process try: - node = nodes[process.pid] + node = nodes[real_process.pid] except KeyError: - node = nodes[process.pid] = cls(process) + node = nodes[real_process.pid] = cls(real_process) finally: try: node.devices.add(process.device) @@ -195,37 +221,40 @@ class TreeNode: # pylint: disable=too-many-instance-attributes return sorted(filter(lambda node: node.is_root, nodes.values()), key=lambda node: node.pid) @staticmethod - def freeze(roots): + def freeze(roots: list[TreeNode]) -> list[FreezedTreeNode]: for root in roots: root.as_snapshot() root.set_prefix() - - return roots - - @staticmethod - def flatten(roots): - flattened = [] - stack = list(reversed(roots)) - while len(stack) > 0: - top = stack.pop() - flattened.append(top) - stack.extend(reversed(top.children)) - return flattened + return roots # type: ignore[return-value] -class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attributes - NAME = 'treeview' - SNAPSHOT_INTERVAL = 0.5 +class FreezedTreeNode(TreeNode): + process: Snapshot # type: ignore[assignment] - def __init__(self, win, root): + +def flatten_process_trees(roots: list[FreezedTreeNode]) -> list[FreezedTreeNode]: + flattened = [] + stack = list(reversed(roots)) + while len(stack) > 0: + top = stack.pop() + flattened.append(top) + stack.extend(reversed(top.children)) + return flattened + + +class TreeViewScreen(BaseSelectableScreen): # pylint: disable=too-many-instance-attributes + NAME: ClassVar[str] = 'treeview' + SNAPSHOT_INTERVAL: ClassVar[float] = 0.5 + + def __init__(self, *, win: curses.window, root: TUI) -> None: super().__init__(win, root) - self.selection = Selection(panel=self) - self.x_offset = 0 - self.y_mouse = None + self.selection: Selection = Selection(self) + self.x_offset: int = 0 + self.y_mouse: int | None = None - self._snapshot_buffer = [] - self._snapshots = [] + self._snapshot_buffer: list[Snapshot] = [] + self._snapshots: list[Snapshot] = [] self.snapshot_lock = threading.Lock() self._snapshot_daemon = threading.Thread( name='treeview-snapshot-daemon', @@ -235,19 +264,19 @@ class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attribut self._daemon_running = threading.Event() self.x, self.y = root.x, root.y - self.scroll_offset = 0 self.width, self.height = root.width, root.height + self.scroll_offset: int = 0 @property - def display_height(self): + def display_height(self) -> int: return self.height - 1 @property - def visible(self): + def visible(self) -> bool: return self._visible @visible.setter - def visible(self, value): + def visible(self, value: bool) -> None: if self._visible != value: self.need_redraw = True self._visible = value @@ -263,11 +292,11 @@ class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attribut self.focused = False @property - def snapshots(self): + def snapshots(self) -> list[Snapshot]: return self._snapshots @snapshots.setter - def snapshots(self, snapshots): + def snapshots(self, snapshots: list[Snapshot]) -> None: with self.snapshot_lock: self.need_redraw = self.need_redraw or len(self._snapshots) > len(snapshots) self._snapshots = snapshots @@ -278,42 +307,46 @@ class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attribut for i, process in enumerate(snapshots): if process._ident[:2] == identity[:2]: # pylint: disable=protected-access self.selection.index = i - self.selection.process = process + self.selection.process = process # type: ignore[assignment] break @classmethod - def set_snapshot_interval(cls, interval): + def set_snapshot_interval(cls, interval: float) -> None: assert interval > 0.0 interval = float(interval) cls.SNAPSHOT_INTERVAL = min(interval / 3.0, 1.0) - cls.take_snapshots = ttl_cache(ttl=interval)( - cls.take_snapshots.__wrapped__, # pylint: disable=no-member + cls.take_snapshots = ttl_cache(ttl=interval)( # type: ignore[method-assign] + cls.take_snapshots.__wrapped__, # type: ignore[attr-defined] # pylint: disable=no-member ) @ttl_cache(ttl=2.0) - def take_snapshots(self): + def take_snapshots(self) -> list[Snapshot]: self.root.main_screen.process_panel.ensure_snapshots() snapshots = ( self.root.main_screen.process_panel._snapshot_buffer # pylint: disable=protected-access ) - roots = TreeNode.merge(snapshots) + roots = TreeNode.merge(snapshots) # type: ignore[arg-type] roots = TreeNode.freeze(roots) - nodes = TreeNode.flatten(roots) + nodes = flatten_process_trees(roots) snapshots = [] for node in nodes: snapshot = node.process snapshot.username = WideString(snapshot.username) snapshot.prefix = node.prefix - if len(node.devices) > 0: - snapshot.devices = 'GPU ' + ','.join( - dev.display_index - for dev in sorted(node.devices, key=lambda device: device.tuple_index) + snapshot.devices = ( + ( + 'GPU ' + + ','.join( + dev.display_index + for dev in sorted(node.devices, key=lambda device: device.tuple_index) + ) ) - else: - snapshot.devices = 'Host' + if node.devices + else 'Host' + ) snapshots.append(snapshot) with self.snapshot_lock: @@ -321,13 +354,13 @@ class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attribut return snapshots - def _snapshot_target(self): + def _snapshot_target(self) -> None: while True: self._daemon_running.wait() self.take_snapshots() time.sleep(self.SNAPSHOT_INTERVAL) - def update_size(self, termsize=None): + def update_size(self, termsize: tuple[int, int] | None = None) -> tuple[int, int]: n_term_lines, n_term_cols = termsize = super().update_size(termsize=termsize) self.width = n_term_cols - self.x @@ -335,7 +368,7 @@ class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attribut return termsize - def poke(self): + def poke(self) -> None: if self._daemon_running.is_set(): self.snapshots = self._snapshot_buffer @@ -362,7 +395,7 @@ class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attribut super().poke() - def draw(self): # pylint: disable=too-many-statements,too-many-locals + def draw(self) -> None: # pylint: disable=too-many-statements,too-many-locals self.color_reset() pid_width = max(3, max((len(str(process.pid)) for process in self.snapshots), default=3)) @@ -407,7 +440,7 @@ class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attribut self.addstr( self.y + 1, self.x, - 'No running GPU processes found{}.'.format(' (in WSL)' if host.WSL else ''), + 'No running GPU processes found{}.'.format(' (in WSL)' if IS_WSL else ''), ) return @@ -450,10 +483,10 @@ class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attribut ) if y == self.y_mouse: - self.selection.process = process + self.selection.process = process # type: ignore[assignment] hint = True - owned = str(process.username) == USERNAME or SUPERUSER + owned = IS_SUPERUSER or str(process.username) == USERNAME if self.selection.is_same_on_host(process): self.color_at( y, @@ -532,19 +565,19 @@ class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attribut attr='bold | reverse', ) - def finalize(self): + def finalize(self) -> None: self.y_mouse = None super().finalize() - def destroy(self): + def destroy(self) -> None: super().destroy() self._daemon_running.clear() - def press(self, key): + def press(self, key: int) -> bool: self.root.keymaps.use_keymap('treeview') - self.root.press(key) + return self.root.press(key) - def click(self, event): + def click(self, event: MouseEvent) -> bool: if event.pressed(1) or event.pressed(3) or event.clicked(1) or event.clicked(3): self.y_mouse = event.y return True @@ -556,51 +589,75 @@ class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attribut self.selection.move(direction=direction) return True - def init_keybindings(self): - def tree_left(): + def init_keybindings(self) -> None: + def tree_left() -> None: self.x_offset = max(0, self.x_offset - 5) - def tree_right(): + def tree_right() -> None: self.x_offset += 5 - def tree_begin(): + def tree_begin() -> None: self.x_offset = 0 - def select_move(direction): + def select_move(direction: int) -> None: self.selection.move(direction=direction) - def select_clear(): + def select_clear() -> None: self.selection.clear() - def tag(): + def tag() -> None: self.selection.tag() select_move(direction=+1) keymaps = self.root.keymaps keymaps.bind('treeview', '', tree_left) - keymaps.copy('treeview', '', '') + keymaps.alias('treeview', '', '') keymaps.bind('treeview', '', tree_right) - keymaps.copy('treeview', '', '') + keymaps.alias('treeview', '', '') keymaps.bind('treeview', '', tree_begin) - keymaps.copy('treeview', '', '^') + keymaps.alias('treeview', '', '^') keymaps.bind('treeview', '', partial(select_move, direction=-1)) - keymaps.copy('treeview', '', '') - keymaps.copy('treeview', '', '') - keymaps.copy('treeview', '', '') - keymaps.copy('treeview', '', '[') + keymaps.alias('treeview', '', '') + keymaps.alias('treeview', '', '') + keymaps.alias('treeview', '', '') + keymaps.alias('treeview', '', '[') keymaps.bind('treeview', '', partial(select_move, direction=+1)) - keymaps.copy('treeview', '', '') - keymaps.copy('treeview', '', '') - keymaps.copy('treeview', '', '') - keymaps.copy('treeview', '', ']') + keymaps.alias('treeview', '', '') + keymaps.alias('treeview', '', '') + keymaps.alias('treeview', '', '') + keymaps.alias('treeview', '', ']') keymaps.bind('treeview', '', partial(select_move, direction=-(1 << 20))) keymaps.bind('treeview', '', partial(select_move, direction=+(1 << 20))) keymaps.bind('treeview', '', select_clear) keymaps.bind('treeview', '', tag) - keymaps.bind('treeview', 'T', partial(send_signal, signal='terminate', panel=self)) - keymaps.bind('treeview', 'K', partial(send_signal, signal='kill', panel=self)) - keymaps.copy('treeview', 'K', 'k') - keymaps.bind('treeview', '', partial(send_signal, signal='interrupt', panel=self)) - keymaps.copy('treeview', '', 'I') + keymaps.bind( + 'treeview', + 'T', + partial( + MessageBox.confirm_sending_signal_to_processes, + signal='terminate', + screen=self, + ), + ) + keymaps.bind( + 'treeview', + 'K', + partial( + MessageBox.confirm_sending_signal_to_processes, + signal='kill', + screen=self, + ), + ) + keymaps.alias('treeview', 'K', 'k') + keymaps.bind( + 'treeview', + '', + partial( + MessageBox.confirm_sending_signal_to_processes, + signal='interrupt', + screen=self, + ), + ) + keymaps.alias('treeview', '', 'I') diff --git a/nvitop/tui/tui.py b/nvitop/tui/tui.py index bce13b2..932e1d0 100644 --- a/nvitop/tui/tui.py +++ b/nvitop/tui/tui.py @@ -3,12 +3,25 @@ # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring +from __future__ import annotations + import curses import shutil import time +from typing import TYPE_CHECKING, Literal, Union -from nvitop.tui.library import ALT_KEY, DisplayableContainer, KeyBuffer, KeyMaps, MouseEvent +from nvitop.tui.library import ( + ALT_KEY, + Device, + DisplayableContainer, + KeyBuffer, + KeyMaps, + MessageBox, + MouseEvent, + Snapshot, +) from nvitop.tui.screens import ( + BaseScreen, BreakLoop, EnvironScreen, HelpScreen, @@ -18,33 +31,44 @@ from nvitop.tui.screens import ( ) -class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + + +__all__ = ['TUI', 'MonitorMode'] + + +MonitorMode = Literal['auto', 'full', 'compact'] + + +class TUI(DisplayableContainer[Union[BaseScreen, MessageBox]]): # pylint: disable=too-many-instance-attributes # pylint: disable-next=too-many-arguments def __init__( self, - devices, - filters=(), + devices: list[Device], + filters: Iterable[Callable[[Snapshot], bool]] = (), *, - ascii=False, # pylint: disable=redefined-builtin - mode='auto', - interval=None, - win=None, - ): + no_unicode: bool = False, + mode: MonitorMode = 'auto', + interval: float | None = None, + win: curses.window | None = None, + ) -> None: super().__init__(win, root=self) self.x = self.y = 0 - self.width = max(79, shutil.get_terminal_size(fallback=(79, 24)).columns - self.x) - self.termsize = None + self.width: int = max(79, shutil.get_terminal_size(fallback=(79, 24)).columns - self.x) + self.height: int = 0 + self.termsize: tuple[int, int] | None = None - self.ascii = ascii + self.no_unicode: bool = no_unicode - self.devices = devices - self.device_count = len(self.devices) + self.devices: list[Device] = devices + self.device_count: int = len(self.devices) - self.main_screen = MainScreen( + self.main_screen: MainScreen = MainScreen( self.devices, filters, - ascii=ascii, + no_unicode=no_unicode, mode=mode, win=win, root=self, @@ -52,29 +76,33 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes self.main_screen.visible = True self.main_screen.focused = False self.add_child(self.main_screen) - self.current_screen = self.previous_screen = self.main_screen + self.current_screen: BaseScreen = self.main_screen + self.previous_screen: BaseScreen = self.main_screen - self._messagebox = None + self._messagebox: MessageBox | None = None if win is not None: - self.environ_screen = EnvironScreen(win=win, root=self) + self.environ_screen: EnvironScreen = EnvironScreen(win=win, root=self) self.environ_screen.visible = False - self.environ_screen.ascii = False + self.environ_screen.no_unicode = False self.add_child(self.environ_screen) - self.treeview_screen = TreeViewScreen(win=win, root=self) + self.treeview_screen: TreeViewScreen = TreeViewScreen(win=win, root=self) self.treeview_screen.visible = False - self.treeview_screen.ascii = self.ascii + self.treeview_screen.no_unicode = self.no_unicode self.add_child(self.treeview_screen) - self.process_metrics_screen = ProcessMetricsScreen(win=win, root=self) + self.process_metrics_screen: ProcessMetricsScreen = ProcessMetricsScreen( + win=win, + root=self, + ) self.process_metrics_screen.visible = False - self.process_metrics_screen.ascii = self.ascii + self.process_metrics_screen.no_unicode = self.no_unicode self.add_child(self.process_metrics_screen) - self.help_screen = HelpScreen(win=win, root=self) + self.help_screen: HelpScreen = HelpScreen(win=win, root=self) self.help_screen.visible = False - self.help_screen.ascii = False + self.help_screen.no_unicode = False self.add_child(self.help_screen) if interval is not None: @@ -86,35 +114,35 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes self.main_screen.process_panel.set_snapshot_interval(interval) self.treeview_screen.set_snapshot_interval(interval) - self.keybuffer = KeyBuffer() - self.keymaps = KeyMaps(self.keybuffer) - self.last_input_time = time.monotonic() + self.keybuffer: KeyBuffer = KeyBuffer() + self.keymaps: KeyMaps = KeyMaps(self.keybuffer) + self.last_input_time: float = time.monotonic() self.init_keybindings() @property - def messagebox(self): + def messagebox(self) -> MessageBox | None: return self._messagebox @messagebox.setter - def messagebox(self, value): + def messagebox(self, value: MessageBox | None) -> None: self.need_redraw = True if self._messagebox is not None: self.remove_child(self._messagebox) self._messagebox = value if value is not None: - self._messagebox.visible = True - self._messagebox.focused = True - self._messagebox.ascii = self.ascii - self._messagebox.previous_focused = self.get_focused_obj() - self.add_child(self._messagebox) + value.visible = True + value.focused = True + value.no_unicode = self.no_unicode + value.previous_focused = self.get_focused_obj() # type: ignore[assignment] + self.add_child(value) - def get_focused_obj(self): + def get_focused_obj(self) -> BaseScreen | MessageBox | None: if self.messagebox is not None: return self.messagebox return super().get_focused_obj() - def update_size(self, termsize=None): + def update_size(self, termsize: tuple[int, int] | None = None) -> tuple[int, int]: n_term_lines, n_term_cols = termsize = super().update_size(termsize=termsize) self.width = n_term_cols - self.x @@ -130,13 +158,14 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes return termsize - def poke(self): + def poke(self) -> None: super().poke() if self.termsize is None: self.update_size() - def draw(self): + def draw(self) -> None: + assert self.win is not None if self.need_redraw: self.win.erase() @@ -151,6 +180,7 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes if not self.need_redraw: return + assert self.termsize is not None n_term_lines, n_term_cols = self.termsize message = ( f'nvitop needs at least a width of 79 to render, the current width is {self.width}.' @@ -178,16 +208,17 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes for y, line in enumerate(lines, start=y_start): self.addstr(y, x_start, line) - def finalize(self): + def finalize(self) -> None: + assert self.win is not None super().finalize() self.win.refresh() - def redraw(self): + def redraw(self) -> None: self.poke() self.draw() self.finalize() - def loop(self): + def loop(self) -> None: if self.win is None: return @@ -200,18 +231,18 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes except BreakLoop: pass - def print(self): + def print(self) -> None: self.main_screen.print() - def handle_mouse(self): + def handle_mouse(self) -> None: """Handle mouse input.""" try: - event = MouseEvent(curses.getmouse()) + event = MouseEvent.get() except curses.error: return super().click(event) - def handle_key(self, key): + def handle_key(self, key: int) -> None: """Handle key input.""" if key < 0: self.keybuffer.clear() @@ -219,11 +250,11 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes self.keymaps.use_keymap('main') self.press(key) - def handle_keys(self, *keys): + def handle_keys(self, *keys: int) -> None: for key in keys: self.handle_key(key) - def press(self, key): + def press(self, key: int) -> bool: keybuffer = self.keybuffer keybuffer.add(key) @@ -238,7 +269,8 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes return False return True - def handle_input(self): # pylint: disable=too-many-branches + def handle_input(self) -> None: # pylint: disable=too-many-branches + assert self.win is not None key = self.win.getch() if key == curses.ERR: return @@ -269,14 +301,14 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes else: self.handle_key(key) - def init_keybindings(self): + def init_keybindings(self) -> None: # pylint: disable=multiple-statements,too-many-statements for screen in self.container: if hasattr(screen, 'init_keybindings'): screen.init_keybindings() - def show_screen(screen, focused=None): + def show_screen(screen: BaseScreen, focused: bool | None = None) -> None: for s in self.container: if s is screen: s.visible = True @@ -288,21 +320,23 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes self.previous_screen = self.current_screen self.current_screen = screen - def show_main(): - show_screen(self.main_screen, focused=False) + def show_main() -> None: + target_screen = self.main_screen + show_screen(target_screen, focused=False) if self.treeview_screen.selection.is_set(): - self.main_screen.selection.process = self.treeview_screen.selection.process + target_screen.selection.process = self.treeview_screen.selection.process self.treeview_screen.selection.clear() self.process_metrics_screen.disable() - def show_environ(): - show_screen(self.environ_screen, focused=True) + def show_environ() -> None: + target_screen = self.environ_screen + show_screen(target_screen, focused=True) if self.previous_screen is not self.help_screen: - self.environ_screen.process = self.previous_screen.selection.process + target_screen.process = self.previous_screen.selection.process # type: ignore[attr-defined] - def environ_return(): + def environ_return() -> None: if self.previous_screen is self.treeview_screen: show_treeview() elif self.previous_screen is self.process_metrics_screen: @@ -310,28 +344,30 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes else: show_main() - def show_treeview(): + def show_treeview() -> None: if not self.main_screen.process_panel.has_snapshots: return - show_screen(self.treeview_screen, focused=True) + target_screen = self.treeview_screen + show_screen(target_screen, focused=True) - if not self.treeview_screen.selection.is_set(): - self.treeview_screen.selection.process = self.main_screen.selection.process + if not target_screen.selection.is_set(): + target_screen.selection.process = self.main_screen.selection.process self.main_screen.selection.clear() - def show_process_metrics(): + def show_process_metrics() -> None: + target_screen = self.process_metrics_screen if self.current_screen is self.main_screen: if self.main_screen.selection.is_set(): - show_screen(self.process_metrics_screen, focused=True) - self.process_metrics_screen.process = self.previous_screen.selection.process + show_screen(target_screen, focused=True) + target_screen.process = self.previous_screen.selection.process # type: ignore[attr-defined] elif self.current_screen is not self.treeview_screen: - show_screen(self.process_metrics_screen, focused=True) + show_screen(target_screen, focused=True) - def show_help(): + def show_help() -> None: show_screen(self.help_screen, focused=True) - def help_return(): + def help_return() -> None: if self.previous_screen is self.treeview_screen: show_treeview() elif self.previous_screen is self.environ_screen: @@ -343,26 +379,26 @@ class TUI(DisplayableContainer): # pylint: disable=too-many-instance-attributes self.keymaps.bind('main', 'e', show_environ) self.keymaps.bind('environ', 'e', environ_return) - self.keymaps.copy('environ', 'e', '') - self.keymaps.copy('environ', 'e', 'q') - self.keymaps.copy('environ', 'e', 'Q') + self.keymaps.alias('environ', 'e', '') + self.keymaps.alias('environ', 'e', 'q') + self.keymaps.alias('environ', 'e', 'Q') self.keymaps.bind('main', 't', show_treeview) self.keymaps.bind('treeview', 't', show_main) - self.keymaps.copy('treeview', 't', 'q') - self.keymaps.copy('treeview', 't', 'Q') + self.keymaps.alias('treeview', 't', 'q') + self.keymaps.alias('treeview', 't', 'Q') self.keymaps.bind('treeview', 'e', show_environ) self.keymaps.bind('main', '', show_process_metrics) self.keymaps.bind('process-metrics', '', show_main) - self.keymaps.copy('process-metrics', '', '') - self.keymaps.copy('process-metrics', '', 'q') - self.keymaps.copy('process-metrics', '', 'Q') + self.keymaps.alias('process-metrics', '', '') + self.keymaps.alias('process-metrics', '', 'q') + self.keymaps.alias('process-metrics', '', 'Q') self.keymaps.bind('process-metrics', 'e', show_environ) - for screen in ('main', 'treeview', 'environ', 'process-metrics'): - self.keymaps.bind(screen, 'h', show_help) - self.keymaps.copy(screen, 'h', '?') + for screen_name in ('main', 'treeview', 'environ', 'process-metrics'): + self.keymaps.bind(screen_name, 'h', show_help) + self.keymaps.alias(screen_name, 'h', '?') self.keymaps.bind('help', '', help_return) self.keymaps.bind('help', '', help_return) diff --git a/pyproject.toml b/pyproject.toml index 511bbe4..2a7958e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -90,10 +90,6 @@ warn_unreachable = true warn_unused_configs = true warn_unused_ignores = true -[[tool.mypy.overrides]] -module = ['nvitop.callbacks.*', 'nvitop.tui.*'] -ignore_errors = true - [tool.pylint] main.py-version = "3.8" basic.good-names = ["x", "y", "dx", "dy", "p", "s", "fg", "bg", "n", "ui", "tx", "rx"] @@ -172,9 +168,6 @@ ignore = [ # TRY003: avoid specifying long messages outside the exception class # long messages are necessary for clarity "TRY003", - # RUF022: `__all__` is not ordered according to an "isort-style" sort - # `__all__` contains comments to group names - "RUF022", ] [tool.ruff.lint.per-file-ignores] @@ -183,7 +176,6 @@ ignore = [ ] "setup.py" = [ "D", # pydocstyle - "ANN", # flake8-annotations ] "nvitop/api/lib*.py" = [ "N", # pep8-naming @@ -194,8 +186,6 @@ ignore = [ ] "nvitop/tui/**/*.py" = [ "D", # pydocstyle - "ANN", # flake8-annotations - "RUF012", # mutable-class-default ] "!nvitop/tui/**/*.py" = [ "TID251", # banned-api diff --git a/setup.py b/setup.py index 4a90aa7..97042cb 100755 --- a/setup.py +++ b/setup.py @@ -76,6 +76,7 @@ if __name__ == '__main__': 'lint': [ 'ruff', 'pylint[spelling]', + 'xdoctest', 'mypy', 'typing-extensions', 'pre-commit',