setup.py

#!/usr/bin/env python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

importdistutils.command.clean
importos
importshutil
importsubprocess
importsys

frompathlibimportPath

fromsetuptoolsimportfind_packages, setup

fromtools.setup_helpers.extensionimportget_ext_modules

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


################################################################################
# Parameters parsed from environment
################################################################################
RUN_BUILD_DEP=True
for_, arginenumerate(sys.argv):
ifargin ["clean", "egg_info", "sdist"]:
RUN_BUILD_DEP=False


def_get_version():
withopen(os.path.join(ROOT_DIR, "version.txt")) asf:
version=f.readline().strip()

sha="Unknown"
try:
sha=subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=str(ROOT_DIR)).decode("ascii").strip()
exceptException:
pass

os_build_version=os.getenv("BUILD_VERSION")
ifos_build_version:
version=os_build_version
elifsha!="Unknown":
version+="+"+sha[:7]

returnversion, sha


def_export_version(version, sha):
version_path=ROOT_DIR/"torchdata"/"version.py"
withopen(version_path, "w") asf:
f.write(f"__version__ = '{version}'\n")
f.write(f"git_version = {repr(sha)}\n")


def_get_requirements():
req_list= []
withPath("requirements.txt").open("r") asf:
forlineinf:
req=line.strip()
iflen(req) ==0orreq.startswith("#"):
continue
req_list.append(req)
returnreq_list


# Use new version of torch on main branch
pytorch_package_dep="torch>=2"

requirements=_get_requirements()
requirements.append(pytorch_package_dep)


classclean(distutils.command.clean.clean):
defrun(self):
# Run default behavior first
distutils.command.clean.clean.run(self)

# Remove torchdata extension
defremove_extension(pattern):
forpathin (ROOT_DIR/"torchdata").glob(pattern):
print(f"removing extension '{path}'")
path.unlink()

forextin ["so", "dylib", "pyd"]:
remove_extension("**/*."+ext)

# Remove build directory
build_dirs= [
ROOT_DIR/"build",
 ]
forpathinbuild_dirs:
ifpath.exists():
print(f"removing '{path}' (and everything under it)")
shutil.rmtree(str(path), ignore_errors=True)


if__name__=="__main__":
VERSION, SHA=_get_version()
_export_version(VERSION, SHA)

print("-- Building version "+VERSION)
setup(
# Metadata
name="torchdata",
version=VERSION,
description="Composable data loading modules for PyTorch",
long_description=Path("README.md").read_text(encoding="utf-8"),
long_description_content_type="text/markdown",
url="https://github.com/pytorch/data",
author="PyTorch Team",
author_email="packages@pytorch.org",
license="BSD",
install_requires=requirements,
python_requires=">=3.9",
classifiers=[
"Intended Audience :: Developers",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: BSD License",
"Operating System :: MacOS :: MacOS X",
"Operating System :: Microsoft :: Windows",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Programming Language :: Python :: Implementation :: CPython",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
 ],
# Package Info
packages=find_packages(exclude=["test*", "examples*", "tools*", "build*"]),
zip_safe=False,
# C++ Extension Modules
ext_modules=get_ext_modules(),
cmdclass={"clean": clean},
 )