Let wheel names get the proper platform suffix

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/

setup.py

@@ -170,4 +170,39 @@ 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 get_and_write_version():
+    if 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
+        version = os.getenv("BUILD_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")
+
+    return version
+
+
+setup(
+    version=get_and_write_version(),
+    ext_modules=[fake_extension],
+    cmdclass={"build_ext": CMakeBuild},
+)

src/torchcodec/__init__.py

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

test/test_version.py

@@ -0,0 +1,6 @@
+import torchcodec
+
+
+def test_version():
+    # Basic test to make sure the attribute exists and is not empty
+    assert torchcodec.__version__

version.txt

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