diff --git a/setup.py b/setup.py index 65dae8ce67c4..c6b5e43b6e9c 100644 --- a/setup.py +++ b/setup.py @@ -257,9 +257,9 @@ from collections import defaultdict import setuptools.command.build_ext -import setuptools.command.install import setuptools.command.sdist from setuptools import Extension, find_packages, setup +from setuptools.command.install import install from setuptools.dist import Distribution from tools.build_pytorch_libs import build_pytorch from tools.generate_torch_version import get_torch_version @@ -268,6 +268,22 @@ from tools.setup_helpers.generate_linker_script import gen_linker_script +class SetupToolsInstallOverride(install): + def run(self): + super().run() + if sys.platform == "win32": + self.remove_shebang_from_easy_install_generated_scripts() + + def remove_shebang_from_easy_install_generated_scripts(self): + for script in self.get_outputs(): + if script.endswith("-script.py"): + with open(script, "rb") as f: + lines = f.readlines() + if len(lines) > 0 and lines[0].startswith(b"#!"): + with open(script, "wb") as f: + f.writelines(lines[1:]) + + def _get_package_path(package_name): spec = importlib.util.find_spec(package_name) if spec: @@ -1068,7 +1084,7 @@ def make_relative_rpath_args(path): "bdist_wheel": wheel_concatenate, "build_ext": build_ext, "clean": clean, - "install": install, + "install": SetupToolsInstallOverride, "sdist": sdist, }