chore(setup): put setup function under a if __name__ == '__main__' block

This commit is contained in:
Xuehai Pan 2025-02-22 16:55:16 +08:00
parent 7e56376f17
commit 5def9731a0
2 changed files with 43 additions and 45 deletions

View file

@ -25,7 +25,7 @@ repos:
- id: debug-statements
- id: double-quote-string-fixer
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.9.4
rev: v0.9.7
hooks:
- id: ruff
args: [--fix, --exit-non-zero-on-fix]
@ -36,7 +36,7 @@ repos:
- id: pyupgrade
args: [--py38-plus] # sync with requires-python
- repo: https://github.com/pycqa/flake8
rev: 7.1.1
rev: 7.1.2
hooks:
- id: flake8
additional_dependencies:

View file

@ -18,12 +18,13 @@ 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 typing import TYPE_CHECKING
from setuptools import setup
if TYPE_CHECKING:
from collections.abc import Generator
from types import ModuleType
@ -70,48 +71,45 @@ def vcs_version(name: str, path: Path | str) -> Generator[ModuleType]:
file.write(content)
extra_requirements = {
'lint': [
'ruff',
'pylint[spelling]',
'mypy',
'typing-extensions',
'pre-commit',
],
'cuda10': ['nvidia-ml-py == 11.450.51'],
}
if __name__ == '__main__':
extra_requirements = {
'lint': [
'ruff',
'pylint[spelling]',
'mypy',
'typing-extensions',
'pre-commit',
],
'cuda10': ['nvidia-ml-py == 11.450.51'],
}
with vcs_version(
name='nvitop.version',
path=HERE / 'nvitop' / 'version.py',
) as version:
for pynvml_major in sorted(
{int(pynvml.partition('.')[0]) for pynvml in version.PYNVML_VERSION_CANDIDATES},
):
pynvml_range = [
pynvml
for pynvml in version.PYNVML_VERSION_CANDIDATES
if pynvml.startswith(f'{pynvml_major}.')
]
if len(pynvml_range) == 1:
extra_requirements[f'cuda{pynvml_major}'] = [
f'nvidia-ml-py == {pynvml_range[0]}',
with vcs_version(name='nvitop.version', path=HERE / 'nvitop' / 'version.py') as version:
for pynvml_major in sorted(
{int(pynvml.partition('.')[0]) for pynvml in version.PYNVML_VERSION_CANDIDATES},
):
pynvml_range = [
pynvml
for pynvml in version.PYNVML_VERSION_CANDIDATES
if pynvml.startswith(f'{pynvml_major}.')
]
elif len(pynvml_range) >= 2:
extra_requirements[f'cuda{pynvml_major}'] = [
f'nvidia-ml-py >= {pynvml_range[0]}, <= {pynvml_range[-1]}',
]
extra_requirements.update(
{
# The identifier could not start with numbers, add a prefix `pynvml-`
f'pynvml-{pynvml}': [f'nvidia-ml-py == {pynvml}']
for pynvml in version.PYNVML_VERSION_CANDIDATES
},
)
if len(pynvml_range) == 1:
extra_requirements[f'cuda{pynvml_major}'] = [
f'nvidia-ml-py == {pynvml_range[0]}',
]
elif len(pynvml_range) >= 2:
extra_requirements[f'cuda{pynvml_major}'] = [
f'nvidia-ml-py >= {pynvml_range[0]}, <= {pynvml_range[-1]}',
]
extra_requirements.update(
{
# The identifier could not start with numbers, add a prefix `pynvml-`
f'pynvml-{pynvml}': [f'nvidia-ml-py == {pynvml}']
for pynvml in version.PYNVML_VERSION_CANDIDATES
},
)
setup(
name='nvitop',
version=version.__version__,
extras_require=extra_requirements,
)
setup(
name='nvitop',
version=version.__version__,
extras_require=extra_requirements,
)