setup.py

#!/usr/bin/env python
import distutils.command.clean
import os
import re
import shutil
import subprocess
from pathlib import Path

import torch
from setuptools import find_packages, setup
from tools import setup_helpers

ROOT_DIR = Path(__file__).parent.resolve()


def _run_cmd(cmd):
    try:
        return subprocess.check_output(cmd, cwd=ROOT_DIR, stderr=subprocess.DEVNULL).decode("ascii").strip()
    except Exception:
        return None


def _get_version(sha):
    with open(ROOT_DIR / "version.txt", "r") as f:
        version = f.read().strip()
    if os.getenv("BUILD_VERSION"):
        version = os.getenv("BUILD_VERSION")
    elif sha is not None:
        version += "+" + sha[:7]
    return version


def _make_version_file(version):
    version_path = ROOT_DIR / "src" / "torchaudio" / "version.py"
    with open(version_path, "w") as f:
        f.write(f"__version__ = '{version}'\n")


def _get_pytorch_version():
    if "PYTORCH_VERSION" in os.environ:
        return f"torch=={os.environ['PYTORCH_VERSION']}"
    return "torch"


class clean(distutils.command.clean.clean):
    def run(self):
        # Run default behavior first
        distutils.command.clean.clean.run(self)

        # Remove torchaudio extension
        for path in (ROOT_DIR / "src").glob("**/*.so"):
            print(f"removing '{path}'")
            path.unlink()
        # Remove build directory
        build_dirs = [
            ROOT_DIR / "build",
        ]
        for path in build_dirs:
            if path.exists():
                print(f"removing '{path}' (and everything under it)")
                shutil.rmtree(str(path), ignore_errors=True)


def _make_version_file(version):
    ROCM_PATH = os.getenv('ROCM_PATH')
    dtk_path = ROCM_PATH + '/.info/rocm_version'
    with open(dtk_path, 'r') as file:
        content = file.read().strip()
    dtk_version = content.replace('.', '')

    hcu_version = f"{version}+das.dtk{dtk_version}"

    version_path = ROOT_DIR / "packaging" / "torchaudio" /"version.py"
    with open(version_path, "w") as f:
        f.write(f"__version__ = '{version}'\n")
        f.write(f"__hcu_version__ = '{hcu_version}'\n")
    return hcu_version


def _parse_url(path):
    with open(path, "r") as file_:
        for line in file_:
            match = re.match(r"^\s*URL\s+(https:\/\/.+)$", line)
            if match:
                url = match.group(1)
                yield url


def _fetch_archives(src):
    for dest, url in src:
        if not dest.exists():
            print(f" --- Fetching {os.path.basename(dest)}")
            torch.hub.download_url_to_file(url, dest, progress=False)


import re

def get_version(file_path: str = "version.txt", encoding: str = "utf-8") -> str | None:
    """
    从指定文件中读取版本号（默认文件为 version.txt）
    """
    try:
        with open(file_path, "r", encoding=encoding) as file:
            line = file.readline().strip()
            
            # 正则匹配语义化版本号格式（如 v1.2.3 或 4.5.6-beta）
            if re.match(r"^v?(?:\d+\.){2}\d+(-\w+)?$", line):
                return line
            else:
                print(f"[错误] 无效的版本号格式: {line}")
                return None

    except FileNotFoundError:
        print(f"[错误] 文件不存在: {file_path}")
    except UnicodeDecodeError:
        print(f"[错误] 编码不匹配，请尝试 encoding='gbk'")
    except Exception as e:
        print(f"[错误] 读取失败: {str(e)}")
    
    return None

def _main():
    sha = _run_cmd(["git", "rev-parse", "HEAD"])
    branch = _run_cmd(["git", "rev-parse", "--abbrev-ref", "HEAD"])
    tag = _run_cmd(["git", "describe", "--tags", "--exact-match", "@"])
    print("-- Git branch:", branch)
    print("-- Git SHA:", sha)
    print("-- Git tag:", tag)
    pytorch_package_dep = _get_pytorch_version()
    print("-- PyTorch dependency:", pytorch_package_dep)
    version= get_version()
    print("-- Building version", version)

    dcu_version = _make_version_file(version)

    with open("README.md") as f:
        long_description = f.read()

    setup(
        name="torchaudio",
        version=dcu_version,
        description="An audio package for PyTorch",
        long_description=long_description,
        long_description_content_type="text/markdown",
        url="https://github.com/pytorch/audio",
        author=(
            "Soumith Chintala, David Pollack, Sean Naren, Peter Goldsborough, "
            "Moto Hira, Caroline Chen, Jeff Hwang, Zhaoheng Ni, Xiaohui Zhang"
        ),
        author_email="soumith@pytorch.org",
        maintainer="Moto Hira, Caroline Chen, Jeff Hwang, Zhaoheng Ni, Xiaohui Zhang",
        maintainer_email="moto@meta.com",
        classifiers=[
            "Environment :: Plugins",
            "Intended Audience :: Developers",
            "Intended Audience :: Science/Research",
            "License :: OSI Approved :: BSD License",
            "Operating System :: MacOS :: MacOS X",
            "Operating System :: Microsoft :: Windows",
            "Operating System :: POSIX",
            "Programming Language :: C++",
            "Programming Language :: Python :: 3.8",
            "Programming Language :: Python :: 3.9",
            "Programming Language :: Python :: 3.10",
            "Programming Language :: Python :: 3.11",
            "Programming Language :: Python :: Implementation :: CPython",
            "Topic :: Multimedia :: Sound/Audio",
            "Topic :: Scientific/Engineering :: Artificial Intelligence",
        ],
        packages=find_packages(where="src"),
        package_dir={"": "src"},
        ext_modules=setup_helpers.get_ext_modules(),
        cmdclass={
            "build_ext": setup_helpers.CMakeBuild,
            "clean": clean,
        },
        install_requires=[pytorch_package_dep],
        zip_safe=False,
    )


if __name__ == "__main__":
    _main()