setup.py 3.64 KB
Newer Older
Soumith Chintala's avatar
Soumith Chintala committed
1
#!/usr/bin/env python
2
import os
3
import re
moto's avatar
moto committed
4
import shutil
5
import subprocess
moto's avatar
moto committed
6
from pathlib import Path
Soumith Chintala's avatar
Soumith Chintala committed
7
from setuptools import setup, find_packages
8
import distutils.command.clean
9

moto's avatar
moto committed
10
from tools import setup_helpers
11

moto's avatar
moto committed
12
ROOT_DIR = Path(__file__).parent.resolve()
13
14


15
16
17
18
19
20
21
def _run_cmd(cmd):
    try:
        return subprocess.check_output(cmd, cwd=ROOT_DIR).decode('ascii').strip()
    except Exception:
        return None


22
# Creating the version file
23
version = '0.11.0a0'
24
sha = _run_cmd(['git', 'rev-parse', 'HEAD'])
25

26
27
if os.getenv('BUILD_VERSION'):
    version = os.getenv('BUILD_VERSION')
28
elif sha is not None:
29
30
31
    version += '+' + sha[:7]
print('-- Building version ' + version)

moto's avatar
moto committed
32
version_path = ROOT_DIR / 'torchaudio' / 'version.py'
33
34
with open(version_path, 'w') as f:
    f.write("__version__ = '{}'\n".format(version))
35
    f.write("git_version = {}\n".format(repr(sha or 'Unknown')))
36

37
pytorch_package_version = os.getenv('PYTORCH_VERSION')
38

39
pytorch_package_dep = 'torch'
40
41
if pytorch_package_version is not None:
    pytorch_package_dep += "==" + pytorch_package_version
42

moto's avatar
moto committed
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

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 / 'torchaudio').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)
peterjc123's avatar
peterjc123 committed
61

62

63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def _get_packages():
    exclude = [
        "build*",
        "test*",
        "torchaudio.csrc*",
        "third_party*",
        "tools*",
    ]
    exclude_prototype = False
    branch_name = _run_cmd(['git', 'rev-parse', '--abbrev-ref', 'HEAD'])
    is_on_tag = _run_cmd(['git', 'describe', '--tags', '--exact-match', '@'])
    print('-- On branch:', branch_name)
    print('-- On tag:', is_on_tag)
    if branch_name is not None and branch_name.startswith('release/'):
        exclude_prototype = True
    if is_on_tag is not None and re.match(r'v[\d.]+(-rc\d+)?', is_on_tag):
        exclude_prototype = True
    if exclude_prototype:
        print('Excluding torchaudio.prototype from the package.')
        exclude.append("torchaudio.prototype")
    return find_packages(exclude=exclude)


Soumith Chintala's avatar
Soumith Chintala committed
86
setup(
87
    name="torchaudio",
88
    version=version,
Soumith Chintala's avatar
Soumith Chintala committed
89
90
    description="An audio package for PyTorch",
    url="https://github.com/pytorch/audio",
91
    author="Soumith Chintala, David Pollack, Sean Naren, Peter Goldsborough",
Soumith Chintala's avatar
Soumith Chintala committed
92
    author_email="soumith@pytorch.org",
Hong Xu's avatar
Hong Xu committed
93
94
95
96
97
98
99
100
101
    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++",
102
103
104
        "Programming Language :: Python :: 3.6",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
Eli Uriegas's avatar
Eli Uriegas committed
105
        "Programming Language :: Python :: 3.9",
Hong Xu's avatar
Hong Xu committed
106
107
108
109
        "Programming Language :: Python :: Implementation :: CPython",
        "Topic :: Multimedia :: Sound/Audio",
        "Topic :: Scientific/Engineering :: Artificial Intelligence"
    ],
110
    packages=_get_packages(),
moto's avatar
moto committed
111
112
    ext_modules=setup_helpers.get_ext_modules(),
    cmdclass={
moto's avatar
moto committed
113
        'build_ext': setup_helpers.CMakeBuild,
114
        'clean': clean,
moto's avatar
moto committed
115
    },
116
117
    install_requires=[pytorch_package_dep],
    zip_safe=False,
118
)