From c00b80d3cf123e001bce263f84cdc8d9e95ba324 Mon Sep 17 00:00:00 2001
From: Xuehai Pan <XuehaiPan@pku.edu.cn>
Date: Fri, 12 Jul 2024 23:07:52 +0800
Subject: [PATCH] feat(api): handle exceptions for function `getpass.getuser()`
 (#130)

---
 .pre-commit-config.yaml           |  2 +-
 CHANGELOG.md                      |  2 +-
 docs/source/spelling_wordlist.txt |  1 +
 nvitop/api/host.py                | 22 +++++++++++++++++++++-
 nvitop/gui/library/utils.py       | 29 +++++++++++++----------------
 nvitop/select.py                  | 12 ++++++------
 6 files changed, 43 insertions(+), 25 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 2548f51..8ac1e60 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -25,7 +25,7 @@ repos:
       - id: debug-statements
       - id: double-quote-string-fixer
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.5.0
+    rev: v0.5.1
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 81058ed..3ca53dd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,7 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
--
+- Handle exceptions for function `getpass.getuser()` by [@XuehaiPan](https://github.com/XuehaiPan) in [#130](https://github.com/XuehaiPan/nvitop/pull/130). Issued by [@landgraf](https://github.com/landgraf).
 
 ### Changed
 
diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt
index af248a4..cb70c1e 100644
--- a/docs/source/spelling_wordlist.txt
+++ b/docs/source/spelling_wordlist.txt
@@ -154,3 +154,4 @@ pytorch
 api
 utils
 GpuStatsLogger
+hostname
diff --git a/nvitop/api/host.py b/nvitop/api/host.py
index 2f75929..0f3ad4b 100644
--- a/nvitop/api/host.py
+++ b/nvitop/api/host.py
@@ -23,7 +23,6 @@ utilization (CPU, memory, disks, network, sensors) in Python.
 from __future__ import annotations
 
 import os as _os
-import time as _time
 from typing import Callable as _Callable
 
 import psutil as _psutil
@@ -31,6 +30,8 @@ from psutil import *  # noqa: F403 # pylint: disable=wildcard-import,unused-wild
 
 
 __all__ = [name for name in _psutil.__all__ if not name.startswith('_')] + [
+    'getuser',
+    'hostname',
     'load_average',
     'uptime',
     'memory_percent',
@@ -60,6 +61,23 @@ LINUX = _psutil.LINUX
 MACOS = _psutil.MACOS
 
 
+def getuser() -> str:
+    """Get the current username from the environment or password database."""
+    import getpass  # pylint: disable=import-outside-toplevel
+
+    try:
+        return getpass.getuser()
+    except (ModuleNotFoundError, OSError):
+        return _os.getlogin()
+
+
+def hostname() -> str:
+    """Get the hostname of the machine."""
+    import platform  # pylint: disable=import-outside-toplevel
+
+    return platform.node()
+
+
 if hasattr(_psutil, 'getloadavg'):
 
     def load_average() -> tuple[float, float, float]:
@@ -75,6 +93,8 @@ else:
 
 def uptime() -> float:
     """Get the system uptime."""
+    import time as _time  # pylint: disable=import-outside-toplevel
+
     return _time.time() - _psutil.boot_time()
 
 
diff --git a/nvitop/gui/library/utils.py b/nvitop/gui/library/utils.py
index 8d12cca..b07ccf5 100644
--- a/nvitop/gui/library/utils.py
+++ b/nvitop/gui/library/utils.py
@@ -3,10 +3,9 @@
 
 # pylint: disable=missing-module-docstring,missing-function-docstring
 
-import getpass
+import contextlib
 import math
 import os
-import platform
 
 from nvitop.api import NA, colored, host, set_color  # noqa: F401 # pylint: disable=unused-import
 from nvitop.gui.library.widestring import WideString
@@ -56,25 +55,23 @@ def make_bar(prefix, percent, width, *, extra_text=''):
     return f'{bar} {text}'.ljust(width)
 
 
-try:
-    USERNAME = getpass.getuser()
-except ModuleNotFoundError:
-    USERNAME = os.getlogin()
+USERNAME = 'N/A'
+with contextlib.suppress(ImportError, OSError):
+    USERNAME = host.getuser()
 
-if host.WINDOWS:
-    import ctypes
+SUPERUSER = False
+with contextlib.suppress(AttributeError, OSError):
+    if host.WINDOWS:
+        import ctypes
 
-    SUPERUSER = bool(ctypes.windll.shell32.IsUserAnAdmin())
-else:
-    try:
-        SUPERUSER = os.geteuid() == 0
-    except AttributeError:
+        SUPERUSER = bool(ctypes.windll.shell32.IsUserAnAdmin())
+    else:
         try:
-            SUPERUSER = os.getuid() == 0
+            SUPERUSER = os.geteuid() == 0
         except AttributeError:
-            SUPERUSER = False
+            SUPERUSER = os.getuid() == 0
 
-HOSTNAME = platform.node()
+HOSTNAME = host.hostname()
 if host.WSL:
     HOSTNAME = f'{HOSTNAME} (WSL)'
 
diff --git a/nvitop/select.py b/nvitop/select.py
index 3e115f0..902bf2a 100644
--- a/nvitop/select.py
+++ b/nvitop/select.py
@@ -57,14 +57,14 @@ Python API:
 from __future__ import annotations
 
 import argparse
-import getpass
+import contextlib
 import math
 import os
 import sys
 import warnings
 from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence, overload
 
-from nvitop.api import Device, GpuProcess, Snapshot, colored, human2bytes, libnvml
+from nvitop.api import Device, GpuProcess, Snapshot, colored, host, human2bytes, libnvml
 from nvitop.version import __version__
 
 
@@ -74,10 +74,10 @@ if TYPE_CHECKING:
 
 __all__ = ['select_devices']
 
-try:
-    USERNAME = getpass.getuser()
-except ModuleNotFoundError:
-    USERNAME = os.getlogin()
+
+USERNAME = 'N/A'
+with contextlib.suppress(ImportError, OSError):
+    USERNAME = host.getuser()
 
 TTY = sys.stdout.isatty()