Let wheel names get the proper platform suffix (#366)

Author: NicolasHug

.gitignore

@@ -10,6 +10,8 @@ src/TorchCodec.egg-info/
 *~
 frame180.* # output from smoke test
 
+src/torchcodec/version.py
+
 docs/build
 # sphinx-gallery
 docs/source/generated_examples/

pyproject.toml

@@ -14,7 +14,7 @@ GitHub = "https://github.com/pytorch/torchcodec"
 Documentation = "https://pytorch.org/torchcodec/stable/index.html"
 
 [tool.setuptools.dynamic]
-version = {attr = "torchcodec.__version__"}
+version = {file = "version.txt"}
 
 [build-system]
 requires = ["setuptools>=61.0"]

setup.py

@@ -170,4 +170,38 @@ def copy_extensions_to_source(self):
 
 # See `CMakeBuild.build_extension()`.
 fake_extension = Extension(name="FAKE_NAME", sources=[])
-setup(ext_modules=[fake_extension], cmdclass={"build_ext": CMakeBuild})
+
+
+def _write_version_files():
+    if version := os.getenv("BUILD_VERSION"):
+        # BUILD_VERSION is set by the `test-infra` build jobs. It typically is
+        # the content of `version.txt` plus some suffix like "+cpu" or "+cu112".
+        # See
+        # https://github.com/pytorch/test-infra/blob/61e6da7a6557152eb9879e461a26ad667c15f0fd/tools/pkg-helpers/pytorch_pkg_helpers/version.py#L113
+        with open(_ROOT_DIR / "version.txt", "w") as f:
+            f.write(f"{version}")
+    else:
+        with open(_ROOT_DIR / "version.txt") as f:
+            version = f.readline().strip()
+        try:
+            sha = (
+                subprocess.check_output(
+                    ["git", "rev-parse", "HEAD"], cwd=str(_ROOT_DIR)
+                )
+                .decode("ascii")
+                .strip()
+            )
+            version += "+" + sha[:7]
+        except Exception:
+            print("INFO: Didn't find sha. Is this a git repo?")
+
+    with open(_ROOT_DIR / "src/torchcodec/version.py", "w") as f:
+        f.write(f"__version__ = '{version}'\n")
+
+
+_write_version_files()
+
+setup(
+    ext_modules=[fake_extension],
+    cmdclass={"build_ext": CMakeBuild},
+)

src/torchcodec/__init__.py

@@ -9,4 +9,7 @@
 from ._frame import Frame, FrameBatch  # usort:skip # noqa
 from . import decoders, samplers  # noqa
 
-__version__ = "0.0.4.dev"
+try:
+    from .version import __version__  # noqa: F401
+except Exception:
+    pass

test/decoders/manual_smoke_test.py

@@ -9,6 +9,8 @@
 import torchcodec
 from torchvision.io.image import write_png
 
+print(f"{torchcodec.__version__ = }")
+
 decoder = torchcodec.decoders._core.create_from_file(
     str(Path(__file__).parent / "../resources/nasa_13013.mp4")
 )

test/test_version.py

@@ -0,0 +1,5 @@
+import torchcodec
+
+
+def test_version():
+    assert torchcodec.__version__

version.txt

@@ -0,0 +1 @@
+0.0.4a0