refactor(setup): refactor setup scripts

XuehaiPan · Jul 4, 2024 · 87ab233 · 87ab233
1 parent 4833f20
commit 87ab233
Show file tree

Hide file tree

Showing 7 changed files with 106 additions and 51 deletions.
diff --git 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.4.7
+    rev: v0.5.0
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
@@ -38,12 +38,12 @@ repos:
     hooks:
       - id: black
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.15.2
+    rev: v3.16.0
     hooks:
       - id: pyupgrade
         args: [--py37-plus] # sync with requires-python
   - repo: https://github.com/pycqa/flake8
-    rev: 7.0.0
+    rev: 7.1.0
     hooks:
       - id: flake8
         additional_dependencies:

diff --git a/nvitop-exporter/nvitop_exporter/version.py b/nvitop-exporter/nvitop_exporter/version.py
@@ -28,8 +28,8 @@
 
     try:
         prefix, sep, suffix = (
-            subprocess.check_output(
-                ['git', 'describe', '--abbrev=7'],  # noqa: S603,S607
+            subprocess.check_output(  # noqa: S603
+                ['git', 'describe', '--abbrev=7'],  # noqa: S607
                 cwd=os.path.dirname(os.path.abspath(__file__)),
                 stderr=subprocess.DEVNULL,
                 text=True,
@@ -43,7 +43,7 @@
         if sep:
             version_prefix, dot, version_tail = prefix.rpartition('.')
             prefix = f'{version_prefix}{dot}{int(version_tail) + 1}'
-            __version__ = sep.join((prefix, suffix))
+            __version__ = f'{prefix}{sep}{suffix}'
             del version_prefix, dot, version_tail
         else:
             __version__ = prefix

diff --git a/nvitop-exporter/setup.py b/nvitop-exporter/setup.py
@@ -2,43 +2,70 @@
 
 """Setup script for ``nvitop-exporter``."""
 
-import pathlib
+from __future__ import annotations
+
+import contextlib
 import re
 import sys
+from importlib.util import module_from_spec, spec_from_file_location
+from pathlib import Path
+from typing import TYPE_CHECKING, Generator
 
 from setuptools import setup
 
 
-HERE = pathlib.Path(__file__).absolute().parent
-VERSION_FILE = HERE / 'nvitop_exporter' / 'version.py'
+if TYPE_CHECKING:
+    from types import ModuleType
+
+
+HERE = Path(__file__).absolute().parent
 
-sys.path.insert(0, str(VERSION_FILE.parent))
-# pylint: disable-next=import-error,wrong-import-position
-import version  # noqa
 
+@contextlib.contextmanager
+def vcs_version(name: str, path: Path | str) -> Generator[ModuleType, None, None]:
+    """Context manager to update version string in a version module."""
+    path = Path(path).absolute()
+    assert path.is_file()
+    module_spec = spec_from_file_location(name=name, location=path)
+    assert module_spec is not None
+    assert module_spec.loader is not None
+    module = sys.modules.get(name)
+    if module is None:
+        module = module_from_spec(module_spec)
+        sys.modules[name] = module
+    module_spec.loader.exec_module(module)
 
-VERSION_CONTENT = None
+    if module.__release__:
+        yield module
+        return
 
-try:
-    if not version.__release__:
+    content = None
+    try:
         try:
-            VERSION_CONTENT = VERSION_FILE.read_text(encoding='utf-8')
-            VERSION_FILE.write_text(
+            content = path.read_text(encoding='utf-8')
+            path.write_text(
                 data=re.sub(
                     r"""__version__\s*=\s*('[^']+'|"[^"]+")""",
-                    f'__version__ = {version.__version__!r}',
-                    string=VERSION_CONTENT,
+                    f'__version__ = {module.__version__!r}',
+                    string=content,
                 ),
                 encoding='utf-8',
             )
         except OSError:
-            VERSION_CONTENT = None
+            content = None
+
+        yield module
+    finally:
+        if content is not None:
+            with path.open(mode='wt', encoding='utf-8', newline='') as file:
+                file.write(content)
+
 
+with vcs_version(
+    name='nvitop_exporter.version',
+    path=(HERE / 'nvitop_exporter' / 'version.py'),
+) as version:
     setup(
         name='nvitop-exporter',
         version=version.__version__,
     )
-finally:
-    if VERSION_CONTENT is not None:
-        with VERSION_FILE.open(mode='wt', encoding='utf-8', newline='') as file:
-            file.write(VERSION_CONTENT)
diff --git a/nvitop/api/device.py b/nvitop/api/device.py
@@ -3126,8 +3126,8 @@ def _parse_cuda_visible_devices(  # pylint: disable=too-many-branches,too-many-s
 
     try:
         raw_uuids = (
-            subprocess.check_output(
-                [  # noqa: S603
+            subprocess.check_output(  # noqa: S603
+                [
                     sys.executable,
                     '-c',
                     textwrap.dedent(

diff --git a/nvitop/api/process.py b/nvitop/api/process.py
@@ -103,6 +103,7 @@ def command_join(cmdline: list[str]) -> str:
 
 _RAISE = object()
 _USE_FALLBACK_WHEN_RAISE = threading.local()  # see also `GpuProcess.failsafe`
+_USE_FALLBACK_WHEN_RAISE.value = False
 
 
 def auto_garbage_clean(
@@ -131,7 +132,7 @@ def wrapped(self: GpuProcess, *args: Any, **kwargs: Any) -> Any:
                 except KeyError:
                     pass
                 # See also `GpuProcess.failsafe`
-                if fallback is _RAISE or not getattr(_USE_FALLBACK_WHEN_RAISE, 'value', False):
+                if fallback is _RAISE or not _USE_FALLBACK_WHEN_RAISE.value:
                     raise
                 if isinstance(fallback, tuple):
                     if isinstance(ex, host.AccessDenied) and fallback == ('No Such Process',):
@@ -1049,7 +1050,7 @@ def failsafe(cls) -> Generator[None, None, None]:
         """  # pylint: disable=line-too-long
         global _USE_FALLBACK_WHEN_RAISE  # pylint: disable=global-statement,global-variable-not-assigned
 
-        prev_value = getattr(_USE_FALLBACK_WHEN_RAISE, 'value', False)
+        prev_value = _USE_FALLBACK_WHEN_RAISE.value
         try:
             _USE_FALLBACK_WHEN_RAISE.value = True
             yield

diff --git a/nvitop/version.py b/nvitop/version.py
@@ -28,8 +28,8 @@
 
     try:
         prefix, sep, suffix = (
-            subprocess.check_output(
-                ['git', 'describe', '--abbrev=7'],  # noqa: S603,S607
+            subprocess.check_output(  # noqa: S603
+                ['git', 'describe', '--abbrev=7'],  # noqa: S607
                 cwd=os.path.dirname(os.path.abspath(__file__)),
                 stderr=subprocess.DEVNULL,
                 text=True,
@@ -43,7 +43,7 @@
         if sep:
             version_prefix, dot, version_tail = prefix.rpartition('.')
             prefix = f'{version_prefix}{dot}{int(version_tail) + 1}'
-            __version__ = sep.join((prefix, suffix))
+            __version__ = f'{prefix}{sep}{suffix}'
             del version_prefix, dot, version_tail
         else:
             __version__ = prefix

diff --git a/setup.py b/setup.py
@@ -11,46 +11,77 @@
 
 """Setup script for ``nvitop``."""
 
-import pathlib
+from __future__ import annotations
+
+import contextlib
 import re
 import sys
+from importlib.util import module_from_spec, spec_from_file_location
+from pathlib import Path
+from typing import TYPE_CHECKING, Generator
 
 from setuptools import setup
 
 
-HERE = pathlib.Path(__file__).absolute().parent
-VERSION_FILE = HERE / 'nvitop' / 'version.py'
+if TYPE_CHECKING:
+    from types import ModuleType
+
+
+HERE = Path(__file__).absolute().parent
 
-sys.path.insert(0, str(VERSION_FILE.parent))
-# pylint: disable-next=import-error,wrong-import-position
-import version  # noqa
 
+@contextlib.contextmanager
+def vcs_version(name: str, path: Path | str) -> Generator[ModuleType, None, None]:
+    """Context manager to update version string in a version module."""
+    path = Path(path).absolute()
+    assert path.is_file()
+    module_spec = spec_from_file_location(name=name, location=path)
+    assert module_spec is not None
+    assert module_spec.loader is not None
+    module = sys.modules.get(name)
+    if module is None:
+        module = module_from_spec(module_spec)
+        sys.modules[name] = module
+    module_spec.loader.exec_module(module)
 
-VERSION_CONTENT = None
+    if module.__release__:
+        yield module
+        return
 
-try:
-    if not version.__release__:
+    content = None
+    try:
         try:
-            VERSION_CONTENT = VERSION_FILE.read_text(encoding='utf-8')
-            VERSION_FILE.write_text(
+            content = path.read_text(encoding='utf-8')
+            path.write_text(
                 data=re.sub(
                     r"""__version__\s*=\s*('[^']+'|"[^"]+")""",
-                    f'__version__ = {version.__version__!r}',
-                    string=VERSION_CONTENT,
+                    f'__version__ = {module.__version__!r}',
+                    string=content,
                 ),
                 encoding='utf-8',
             )
         except OSError:
-            VERSION_CONTENT = None
+            content = None
+
+        yield module
+    finally:
+        if content is not None:
+            with path.open(mode='wt', encoding='utf-8', newline='') as file:
+                file.write(content)
+
 
+with vcs_version(
+    name='nvitop.version',
+    path=HERE / 'nvitop' / 'version.py',
+) as version:
     setup(
         name='nvitop',
         version=version.__version__,
         extras_require={
             'lint': [
-                'black >= 22.6.0',
+                'black >= 24.0.0, < 25.0.0a0',
                 'isort',
-                'pylint[spelling] >= 2.16.0',
+                'pylint[spelling]',
                 'mypy',
                 'typing-extensions',
                 'pre-commit',
@@ -63,7 +94,3 @@
             },
         },
     )
-finally:
-    if VERSION_CONTENT is not None:
-        with VERSION_FILE.open(mode='wt', encoding='utf-8', newline='') as file:
-            file.write(VERSION_CONTENT)