github-starred/nvitop
2
0
Fork 0
mirror of https://github.com/XuehaiPan/nvitop.git synced 2026-05-21 06:45:24 -06:00
Code Issues 29 Projects Packages Wiki Activity Actions

lint(tui): add doctests and add type annotations in nvitop.tui (#164)

Browse source
This commit is contained in:
Xuehai Pan 2025-04-25 05:50:41 +08:00 • committed by GitHub
parent dfb4e3bf55
commit 96a0fab34f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
48 changed files with 2302 additions and 1505 deletions
Download patch file Download diff file Expand all files Collapse all files

9
.github/workflows/lint.yaml vendored
View file

@ -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

1
CHANGELOG.md
View file

@ -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

2
docs/source/spelling_wordlist.txt
View file

@ -24,6 +24,7 @@ devicesnapshot
displayables
divmod
docstring
doctest
ecc
enum
env
@ -107,6 +108,7 @@ rtx
runtime
rw
rx
selectable
shader
sm
smi

2
nvitop/api/__init__.py
View file

@ -65,7 +65,7 @@ from nvitop.api.utils import ( # explicitly export these to appease mypy
)
__all__ = [
__all__ = [ # noqa: RUF022
'NVMLError',
'nvmlCheckReturn',
'libnvml',

2
nvitop/api/collector.py
View file

@ -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

35
nvitop/api/device.py
View file

@ -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-<GPU-UUID>`
# MIG UUID : `MIG-GPU-<GPU-UUID>/<GPU instance ID>/<compute instance ID>`
# MIG UUID (R470+): `MIG-<MIG-UUID>`
UUID_PATTERN: re.Pattern = re.compile(
UUID_PATTERN: ClassVar[re.Pattern] = re.compile(
r"""^ # full match
(?:(?P<MigMode>MIG)-)? # prefix for MIG UUID
(?:(?P<GpuUuid>GPU)-)? # 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

8
nvitop/api/host.py
View file

@ -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']

4
nvitop/api/libcuda.py
View file

@ -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

6
nvitop/api/libcudart.py
View file

@ -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

10
nvitop/api/libnvml.py
View file

@ -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]] = {'<default>': '%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:

4
nvitop/api/process.py
View file

@ -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."""

3
nvitop/api/termcolor.py
View file

@ -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',

86
nvitop/api/utils.py
View file

@ -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<size>\d+(?:\.\d+)?)\s*(?P<unit>[KMGTP]i?B?|B?)\s*$',
r'^\s*\+?\s*(?P<size>\d+(?:\.\d+)?)\s*(?P<unit>([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()

2
nvitop/callbacks/keras.py
View file

@ -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

2
nvitop/callbacks/lightning.py
View file

@ -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

2
nvitop/callbacks/pytorch_lightning.py
View file

@ -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

2
nvitop/callbacks/tensorboard.py
View file

@ -17,6 +17,8 @@
# pylint: disable=missing-module-docstring
# mypy: ignore-errors
from __future__ import annotations
import warnings

14
nvitop/cli.py
View file

@ -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:

11
nvitop/select.py
View file

@ -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

14
nvitop/tui/__init__.py
View file

@ -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',
]

59
nvitop/tui/library/__init__.py
View file

@ -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',
]

126
nvitop/tui/library/device.py
View file

@ -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',

107
nvitop/tui/library/displayable.py
View file

@ -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:

206
nvitop/tui/library/history.py
View file

@ -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

56
nvitop/tui/library/host.py
View file

@ -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

250
nvitop/tui/library/keybinding.py
View file

@ -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 <BS> 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 <M-*> 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 <M-*> 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 <M-*> in REVERSED_SPECIAL_KEYS
SPECIAL_KEYS[f'M-{char}'] = (ALT_KEY, ord(char))
SPECIAL_KEYS[f'A-{char}'] = (ALT_KEY, ord(char)) # overrides <M-*> in REVERSED_SPECIAL_KEYS
# We will need to reorder the keys of SPECIAL_KEYS below.
# For example, <C-j> will override <Enter> in REVERSE_SPECIAL_KEYS,
# this makes construct_keybinding(parse_keybinding('<CR>')) == '<C-j>'
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<CR>'))
>>> out = parse_keybinding('lol<CR>')
>>> out
(108, 111, 108, 10)
>>> out == (ord('l'), ord('o'), ord('l'), ord('\n'))
True
>>> out = tuple(parse_keybinding('x<A-Left>'))
>>> out = parse_keybinding('x<A-Left>')
>>> out
(120, 9003, 260)
>>> out == (ord('x'), ALT_KEY, curses.KEY_LEFT)
True
"""
assert isinstance(obj, (tuple, int, str))
def parse(obj: IntKey | str) -> Generator[int]: # pylint: disable=too-many-branches
if isinstance(obj, tuple):
yield from obj
elif isinstance(obj, int): # pylint: disable=too-many-nested-blocks
return
if isinstance(obj, int):
yield obj
else: # pylint: disable=too-many-nested-blocks
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<CR>'))
@ -187,11 +227,7 @@ def construct_keybinding(keys):
>>> construct_keybinding(parse_keybinding('x<Alt><Left>'))
'x<A-Left>'
"""
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<CR>')
@ -225,72 +261,74 @@ 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):
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):
KeyMaps._unbind_traverse(pointer[keys[pos]], keys, pos=pos + 1)
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:
@ -303,35 +341,45 @@ class KeyMaps(dict):
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()

155
nvitop/tui/library/libcurses.py
View file

@ -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):

235
nvitop/tui/library/messagebox.py
View file

@ -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,85 +230,87 @@ 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({'<Enter>', '<Esc>', '<Left>', '<Right>'}.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', '<Esc>', partial(self.apply, index=self.cancel, wait=False))
if 'q' not in keymaps['messagebox'] and 'Q' not in keymaps['messagebox']:
keymaps.copy('messagebox', '<Esc>', 'q')
keymaps.copy('messagebox', '<Esc>', 'Q')
if ord('q') not in keymap and ord('Q') not in keymap:
keymaps.alias('messagebox', '<Esc>', 'q')
keymaps.alias('messagebox', '<Esc>', 'Q')
keymaps.bind('messagebox', '<Enter>', self.apply)
if '<Space>' not in keymaps['messagebox']:
keymaps.copy('messagebox', '<Enter>', '<Space>')
if NAMED_SPECIAL_KEYS['Space'] not in keymap:
keymaps.alias('messagebox', '<Enter>', '<Space>')
keymaps.bind('messagebox', '<Left>', select_previous)
keymaps.bind('messagebox', '<Right>', select_next)
if ',' not in keymaps['messagebox'] and '.' not in keymaps['messagebox']:
keymaps.copy('messagebox', '<Left>', ',')
keymaps.copy('messagebox', '<Right>', '.')
if '<' not in keymaps['messagebox'] and '>' not in keymaps['messagebox']:
keymaps.copy('messagebox', '<Left>', '<')
keymaps.copy('messagebox', '<Right>', '>')
if '[' not in keymaps['messagebox'] and ']' not in keymaps['messagebox']:
keymaps.copy('messagebox', '<Left>', '[')
keymaps.copy('messagebox', '<Right>', ']')
if '<Tab>' not in keymaps['messagebox'] and '<S-Tab>' not in keymaps['messagebox']:
keymaps.copy('messagebox', '<Left>', '<S-Tab>')
keymaps.copy('messagebox', '<Right>', '<Tab>')
if ord(',') not in keymap and ord('.') not in keymap:
keymaps.alias('messagebox', '<Left>', ',')
keymaps.alias('messagebox', '<Right>', '.')
if ord('<') not in keymap and ord('>') not in keymap:
keymaps.alias('messagebox', '<Left>', '<')
keymaps.alias('messagebox', '<Right>', '>')
if ord('[') not in keymap and ord(']') not in keymap:
keymaps.alias('messagebox', '<Left>', '[')
keymaps.alias('messagebox', '<Right>', ']')
if NAMED_SPECIAL_KEYS['Tab'] not in keymap and NAMED_SPECIAL_KEYS['S-Tab'] not in keymap:
keymaps.alias('messagebox', '<Left>', '<S-Tab>')
keymaps.alias('messagebox', '<Right>', '<Tab>')
def send_signal(signal, panel):
assert signal in {'terminate', 'kill', 'interrupt'}
@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 panel.selection.processes():
for process in screen.selection.processes():
try:
username = process.username()
except host.PsutilError:
@ -279,14 +326,13 @@ def send_signal(signal, panel):
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(
screen.root.messagebox = MessageBox(
message=message,
options=[
MessageBox.Option(
'SIGTERM',
't',
panel.selection.terminate,
screen.selection.terminate,
keys=('T',),
attrs=(
{'y': 0, 'x': 0, 'width': 7, 'fg': 'red'},
@ -296,7 +342,7 @@ def send_signal(signal, panel):
MessageBox.Option(
'SIGKILL',
'k',
panel.selection.kill,
screen.selection.kill,
keys=('K',),
attrs=(
{'y': 0, 'x': 0, 'width': 7, 'fg': 'red'},
@ -306,7 +352,7 @@ def send_signal(signal, panel):
MessageBox.Option(
'SIGINT',
'i',
panel.selection.interrupt,
screen.selection.interrupt,
keys=('I',),
attrs=(
{'y': 0, 'x': 0, 'width': 6, 'fg': 'red'},
@ -325,7 +371,6 @@ def send_signal(signal, panel):
yes=None,
no=3,
cancel=3,
win=panel.win,
root=panel.root,
win=screen.win, # type: ignore[arg-type]
root=screen.root,
)
# pylint: enable=use-dict-literal

68
nvitop/tui/library/mouse.py
View file

@ -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

29
nvitop/tui/library/process.py
View file

@ -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

103
nvitop/tui/library/selection.py
View file

@ -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

50
nvitop/tui/library/utils.py
View file

@ -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'

140
nvitop/tui/library/widestring.py
View file

@ -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]
<WideString 'd'>
@ -109,21 +119,21 @@ class WideString: # pylint: disable=too-few-public-methods,wrong-spelling-in-do
<WideString 'sd'>
>>> WideString('asdf')[1:-100]
<WideString ''>
>>> WideString('モヒカン')[2:4]
<WideString 'ヒ'>
>>> WideString('モヒカン')[2:5]
<WideString 'ヒ '>
>>> WideString('モabカン')[2:5]
>>> WideString('十百千万')[2:4]
<WideString '百'>
>>> WideString('十百千万')[2:5]
<WideString '百 '>
>>> WideString('十ab千万')[2:5]
<WideString 'ab '>
>>> WideString('モヒカン')[1:5]
<WideString ' ヒ '>
>>> WideString('モヒカン')[:]
<WideString 'モヒカン'>
>>> WideString('aモ')[0:3]
<WideString 'aモ'>
>>> WideString('aモ')[0:2]
>>> WideString('十百千万')[1:5]
<WideString ' 百 '>
>>> WideString('十百千万')[:]
<WideString '十百千万'>
>>> WideString('a十')[0:3]
<WideString 'a十'>
>>> WideString('a十')[0:2]
<WideString 'a '>
>>> WideString('aモ')[0:1]
>>> WideString('a十')[0:1]
<WideString 'a'>
"""
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'>
>>> WideString('poo').ljust(5)
<WideString 'poo '>
>>> WideString('モヒカン').ljust(10)
<WideString 'モヒカン '>
>>> WideString('十百千万').ljust(10)
<WideString '十百千万 '>
"""
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'>
>>> WideString('poo').rjust(5)
<WideString ' poo'>
>>> WideString('モヒカン').rljust(10)
<WideString ' モヒカン'>
>>> WideString('十百千万').rjust(10)
<WideString ' 十百千万'>
"""
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'>
>>> WideString('poo').center(5)
<WideString ' poo '>
>>> WideString('モヒカン').center(10)
<WideString ' モヒカン '>
>>> WideString('十百千万').center(10)
<WideString ' 十百千万 '>
"""
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 'poo'>
>>> WideString(' モヒカン ').strip()
<WideString 'モヒカン'>
>>> WideString(' 十百千万 ').strip()
<WideString '十百千万'>
"""
return WideString(self.string.strip(chars))
def lstrip(self, chars=None):
def lstrip(self, chars: str | None = None) -> WideString:
"""
>>> WideString(' poo ').lstrip()
<WideString 'poo '>
>>> WideString(' モヒカン ').lstrip()
<WideString 'モヒカン '>
>>> WideString(' 十百千万 ').lstrip()
<WideString '十百千万 '>
"""
return WideString(self.string.lstrip(chars))
def rstrip(self, chars=None):
def rstrip(self, chars: str | None = None) -> WideString:
"""
>>> WideString(' poo ').rstrip()
<WideString ' poo'>
>>> WideString(' モヒカン ').rstrip()
<WideString ' モヒカン'>
>>> WideString(' 十百千万 ').rstrip()
<WideString ' 十百千万'>
"""
return WideString(self.string.rstrip(chars))
if __name__ == '__main__':
import doctest
doctest.testmod()

13
nvitop/tui/screens/__init__.py
View file

@ -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',
]

32
nvitop/tui/screens/base.py Normal file
View file

@ -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

145
nvitop/tui/screens/environ.py
View file

@ -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', '<C-r>')
keymaps.copy('environ', 'r', '<F5>')
keymaps.alias('environ', 'r', 'R')
keymaps.alias('environ', 'r', '<C-r>')
keymaps.alias('environ', 'r', '<F5>')
keymaps.bind('environ', '<Left>', environ_left)
keymaps.copy('environ', '<Left>', '<A-h>')
keymaps.alias('environ', '<Left>', '<A-h>')
keymaps.bind('environ', '<Right>', environ_right)
keymaps.copy('environ', '<Right>', '<A-l>')
keymaps.alias('environ', '<Right>', '<A-l>')
keymaps.bind('environ', '<C-a>', environ_begin)
keymaps.copy('environ', '<C-a>', '^')
keymaps.alias('environ', '<C-a>', '^')
keymaps.bind('environ', '<Up>', partial(environ_move, direction=-1))
keymaps.copy('environ', '<Up>', '<S-Tab>')
keymaps.copy('environ', '<Up>', '<A-k>')
keymaps.copy('environ', '<Up>', '<PageUp>')
keymaps.copy('environ', '<Up>', '[')
keymaps.alias('environ', '<Up>', '<S-Tab>')
keymaps.alias('environ', '<Up>', '<A-k>')
keymaps.alias('environ', '<Up>', '<PageUp>')
keymaps.alias('environ', '<Up>', '[')
keymaps.bind('environ', '<Down>', partial(environ_move, direction=+1))
keymaps.copy('environ', '<Down>', '<Tab>')
keymaps.copy('environ', '<Down>', '<A-j>')
keymaps.copy('environ', '<Down>', '<PageDown>')
keymaps.copy('environ', '<Down>', ']')
keymaps.alias('environ', '<Down>', '<Tab>')
keymaps.alias('environ', '<Down>', '<A-j>')
keymaps.alias('environ', '<Down>', '<PageDown>')
keymaps.alias('environ', '<Down>', ']')
keymaps.bind('environ', '<Home>', partial(environ_move, direction=-(1 << 20)))
keymaps.bind('environ', '<End>', partial(environ_move, direction=+(1 << 20)))

36
nvitop/tui/screens/help.py
View file

@ -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)

203
nvitop/tui/screens/main/__init__.py
View file

@ -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', '<C-r>')
keymaps.copy('main', 'r', '<F5>')
keymaps.alias('main', 'r', 'R')
keymaps.alias('main', 'r', '<C-r>')
keymaps.alias('main', 'r', '<F5>')
keymaps.bind('main', '<PageUp>', partial(screen_move, direction=-1))
keymaps.copy('main', '<PageUp>', '[')
keymaps.copy('main', '<PageUp>', '<A-K>')
keymaps.alias('main', '<PageUp>', '[')
keymaps.alias('main', '<PageUp>', '<A-K>')
keymaps.bind('main', '<PageDown>', partial(screen_move, direction=+1))
keymaps.copy('main', '<PageDown>', ']')
keymaps.copy('main', '<PageDown>', '<A-J>')
keymaps.alias('main', '<PageDown>', ']')
keymaps.alias('main', '<PageDown>', '<A-J>')
keymaps.bind('main', '<Left>', host_left)
keymaps.copy('main', '<Left>', '<A-h>')
keymaps.alias('main', '<Left>', '<A-h>')
keymaps.bind('main', '<Right>', host_right)
keymaps.copy('main', '<Right>', '<A-l>')
keymaps.alias('main', '<Right>', '<A-l>')
keymaps.bind('main', '<C-a>', host_begin)
keymaps.copy('main', '<C-a>', '^')
keymaps.alias('main', '<C-a>', '^')
keymaps.bind('main', '<C-e>', host_end)
keymaps.copy('main', '<C-e>', '$')
keymaps.alias('main', '<C-e>', '$')
keymaps.bind('main', '<Up>', partial(select_move, direction=-1))
keymaps.copy('main', '<Up>', '<S-Tab>')
keymaps.copy('main', '<Up>', '<A-k>')
keymaps.alias('main', '<Up>', '<S-Tab>')
keymaps.alias('main', '<Up>', '<A-k>')
keymaps.bind('main', '<Down>', partial(select_move, direction=+1))
keymaps.copy('main', '<Down>', '<Tab>')
keymaps.copy('main', '<Down>', '<A-j>')
keymaps.alias('main', '<Down>', '<Tab>')
keymaps.alias('main', '<Down>', '<A-j>')
keymaps.bind('main', '<Home>', partial(select_move, direction=-(1 << 20)))
keymaps.bind('main', '<End>', partial(select_move, direction=+(1 << 20)))
keymaps.bind('main', '<Esc>', select_clear)
keymaps.bind('main', '<Space>', 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', '<C-c>', partial(send_signal, signal='interrupt', panel=self))
keymaps.copy('main', '<C-c>', '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',
'<C-c>',
partial(
MessageBox.confirm_sending_signal_to_processes,
signal='interrupt',
screen=self,
),
)
keymaps.alias('main', '<C-c>', '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(

19
nvitop/tui/screens/main/panels/__init__.py Normal file
View file

@ -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',
]

34
nvitop/tui/screens/main/panels/base.py Normal file
View file

@ -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

149
nvitop/tui/screens/main/device.py → nvitop/tui/screens/main/panels/device.py
View file

@ -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)

152
nvitop/tui/screens/main/host.py → nvitop/tui/screens/main/panels/host.py
View file

@ -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
load_average = 'Load Average: {} {} {}'.format(
*(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)
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)

239
nvitop/tui/screens/main/process.py → nvitop/tui/screens/main/panels/process.py
View file

@ -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,18 +64,24 @@ 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 = {
ORDERS: ClassVar[Mapping[OrderName, Order]] = MappingProxyType(
{
'natural': Order(
key=attrgetter('device.tuple_index', '_gone', 'username', 'pid'),
key=attrgetter(
'device.tuple_index',
'_gone',
'username',
'pid',
),
reverse=False,
offset=3,
column='ID',
@ -60,7 +90,11 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes
bind_key='n',
),
'pid': Order(
key=attrgetter('_gone', 'pid', 'device.tuple_index'),
key=attrgetter(
'_gone',
'pid',
'device.tuple_index',
),
reverse=False,
offset=10,
column='PID',
@ -69,7 +103,12 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes
bind_key='p',
),
'username': Order(
key=attrgetter('_gone', 'username', 'pid', 'device.tuple_index'),
key=attrgetter(
'_gone',
'username',
'pid',
'device.tuple_index',
),
reverse=False,
offset=19,
column='USER',
@ -121,12 +160,18 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes
reverse=True,
offset=38,
column='GMBW',
previous='gpu_memory',
previous='sm_utilization',
next='cpu_percent',
bind_key='b',
),
'cpu_percent': Order(
key=attrgetter('_gone', 'cpu_percent', 'memory_percent', 'pid', 'device.tuple_index'),
key=attrgetter(
'_gone',
'cpu_percent',
'memory_percent',
'pid',
'device.tuple_index',
),
reverse=True,
offset=44,
column='%CPU',
@ -135,7 +180,13 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes
bind_key='c',
),
'memory_percent': Order(
key=attrgetter('_gone', 'memory_percent', 'cpu_percent', 'pid', 'device.tuple_index'),
key=attrgetter(
'_gone',
'memory_percent',
'cpu_percent',
'pid',
'device.tuple_index',
),
reverse=True,
offset=50,
column='%MEM',
@ -144,7 +195,12 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes
bind_key='m',
),
'time': Order(
key=attrgetter('_gone', 'running_time', 'pid', 'device.tuple_index'),
key=attrgetter(
'_gone',
'running_time',
'pid',
'device.tuple_index',
),
reverse=True,
offset=56,
column='TIME',
@ -152,32 +208,41 @@ class ProcessPanel(Displayable): # pylint: disable=too-many-instance-attributes
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)

174
nvitop/tui/screens/metrics.py
View file

@ -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)

249
nvitop/tui/screens/treeview.py
View file

@ -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,15 +221,18 @@ 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 # type: ignore[return-value]
return roots
@staticmethod
def flatten(roots):
class FreezedTreeNode(TreeNode):
process: Snapshot # type: ignore[assignment]
def flatten_process_trees(roots: list[FreezedTreeNode]) -> list[FreezedTreeNode]:
flattened = []
stack = list(reversed(roots))
while len(stack) > 0:
@ -213,19 +242,19 @@ class TreeNode: # pylint: disable=too-many-instance-attributes
return flattened
class TreeViewScreen(Displayable): # pylint: disable=too-many-instance-attributes
NAME = 'treeview'
SNAPSHOT_INTERVAL = 0.5
class TreeViewScreen(BaseSelectableScreen): # pylint: disable=too-many-instance-attributes
NAME: ClassVar[str] = 'treeview'
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.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(
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', '<Left>', tree_left)
keymaps.copy('treeview', '<Left>', '<A-h>')
keymaps.alias('treeview', '<Left>', '<A-h>')
keymaps.bind('treeview', '<Right>', tree_right)
keymaps.copy('treeview', '<Right>', '<A-l>')
keymaps.alias('treeview', '<Right>', '<A-l>')
keymaps.bind('treeview', '<C-a>', tree_begin)
keymaps.copy('treeview', '<C-a>', '^')
keymaps.alias('treeview', '<C-a>', '^')
keymaps.bind('treeview', '<Up>', partial(select_move, direction=-1))
keymaps.copy('treeview', '<Up>', '<S-Tab>')
keymaps.copy('treeview', '<Up>', '<A-k>')
keymaps.copy('treeview', '<Up>', '<PageUp>')
keymaps.copy('treeview', '<Up>', '[')
keymaps.alias('treeview', '<Up>', '<S-Tab>')
keymaps.alias('treeview', '<Up>', '<A-k>')
keymaps.alias('treeview', '<Up>', '<PageUp>')
keymaps.alias('treeview', '<Up>', '[')
keymaps.bind('treeview', '<Down>', partial(select_move, direction=+1))
keymaps.copy('treeview', '<Down>', '<Tab>')
keymaps.copy('treeview', '<Down>', '<A-j>')
keymaps.copy('treeview', '<Down>', '<PageDown>')
keymaps.copy('treeview', '<Down>', ']')
keymaps.alias('treeview', '<Down>', '<Tab>')
keymaps.alias('treeview', '<Down>', '<A-j>')
keymaps.alias('treeview', '<Down>', '<PageDown>')
keymaps.alias('treeview', '<Down>', ']')
keymaps.bind('treeview', '<Home>', partial(select_move, direction=-(1 << 20)))
keymaps.bind('treeview', '<End>', partial(select_move, direction=+(1 << 20)))
keymaps.bind('treeview', '<Esc>', select_clear)
keymaps.bind('treeview', '<Space>', 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', '<C-c>', partial(send_signal, signal='interrupt', panel=self))
keymaps.copy('treeview', '<C-c>', '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',
'<C-c>',
partial(
MessageBox.confirm_sending_signal_to_processes,
signal='interrupt',
screen=self,
),
)
keymaps.alias('treeview', '<C-c>', 'I')

196
nvitop/tui/tui.py
View file

@ -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', '<Esc>')
self.keymaps.copy('environ', 'e', 'q')
self.keymaps.copy('environ', 'e', 'Q')
self.keymaps.alias('environ', 'e', '<Esc>')
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', '<Enter>', show_process_metrics)
self.keymaps.bind('process-metrics', '<Enter>', show_main)
self.keymaps.copy('process-metrics', '<Enter>', '<Esc>')
self.keymaps.copy('process-metrics', '<Enter>', 'q')
self.keymaps.copy('process-metrics', '<Enter>', 'Q')
self.keymaps.alias('process-metrics', '<Enter>', '<Esc>')
self.keymaps.alias('process-metrics', '<Enter>', 'q')
self.keymaps.alias('process-metrics', '<Enter>', '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', '<Esc>', help_return)
self.keymaps.bind('help', '<any>', help_return)

10
pyproject.toml
View file

@ -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

1
setup.py
View file

@ -76,6 +76,7 @@ if __name__ == '__main__':
'lint': [
'ruff',
'pylint[spelling]',
'xdoctest',
'mypy',
'typing-extensions',
'pre-commit',