Clean up setup.py (#1900)

setup.py

@@ -19,26 +19,27 @@ def _run_cmd(cmd):
         return None
 
 
-# Creating the version file
-version = '0.11.0a0'
-sha = _run_cmd(['git', 'rev-parse', 'HEAD'])
-
-if os.getenv('BUILD_VERSION'):
+def _get_version(sha):
+    version = '0.11.0a0'
+    if os.getenv('BUILD_VERSION'):
         version = os.getenv('BUILD_VERSION')
-elif sha is not None:
+    elif sha is not None:
         version += '+' + sha[:7]
-print('-- Building version ' + version)
+    return version
+
 
-version_path = ROOT_DIR / 'torchaudio' / 'version.py'
-with open(version_path, 'w') as f:
-    f.write("__version__ = '{}'\n".format(version))
-    f.write("git_version = {}\n".format(repr(sha or 'Unknown')))
+def _make_version_file(version, sha):
+    sha = 'Unknown' if sha is None else sha
+    version_path = ROOT_DIR / 'torchaudio' / 'version.py'
+    with open(version_path, 'w') as f:
+        f.write(f"__version__ = '{version}'\n")
+        f.write(f"git_version = '{sha}'\n")
 
-pytorch_package_version = os.getenv('PYTORCH_VERSION')
 
-pytorch_package_dep = 'torch'
-if pytorch_package_version is not None:
-    pytorch_package_dep += "==" + pytorch_package_version
+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):
@@ -60,7 +61,7 @@ class clean(distutils.command.clean.clean):
                 shutil.rmtree(str(path), ignore_errors=True)
 
 
-def _get_packages():
+def _get_packages(branch_name, tag):
     exclude = [
         "build*",
         "test*",
@@ -69,13 +70,9 @@ def _get_packages():
         "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):
+    if tag is not None and re.match(r'v[\d.]+(-rc\d+)?', tag):
         exclude_prototype = True
     if exclude_prototype:
         print('Excluding torchaudio.prototype from the package.')
@@ -83,7 +80,21 @@ def _get_packages():
     return find_packages(exclude=exclude)
 
 
-setup(
+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(sha)
+    print('-- Building version', version)
+
+    _make_version_file(version, sha)
+
+    setup(
         name="torchaudio",
         version=version,
         description="An audio package for PyTorch",
@@ -107,7 +118,7 @@ setup(
             "Topic :: Multimedia :: Sound/Audio",
             "Topic :: Scientific/Engineering :: Artificial Intelligence"
         ],
-    packages=_get_packages(),
+        packages=_get_packages(branch, tag),
         ext_modules=setup_helpers.get_ext_modules(),
         cmdclass={
             'build_ext': setup_helpers.CMakeBuild,
@@ -115,4 +126,8 @@ setup(
         },
         install_requires=[pytorch_package_dep],
         zip_safe=False,
-)
+    )
+
+
+if __name__ == '__main__':
+    _main()