Add nightly

setup.py

@@ -14,16 +14,32 @@
 # ==============================================================================
 from setuptools import find_packages, setup
 import versioneer
+import os
 
 
 readme = open("README.md").read()
+NAME = "jaxutils"
 
 
-REQUIRES = [
-    "jax>=0.4.0",
-    "jaxlib>=0.4.0",
-    "jaxtyping"
-]
+# Handle builds of nightly release - adapted from BlackJax.
+if "BUILD_JAXUTILS_NIGHTLY" in os.environ:
+    if os.environ["BUILD_JAXUTILS_NIGHTLY"] == "nightly":
+        NAME += "-nightly"
+
+        from versioneer import get_versions as original_get_versions
+
+        def get_versions():
+            from datetime import datetime, timezone
+
+            suffix = datetime.now(timezone.utc).strftime(r".dev%Y%m%d")
+            versions = original_get_versions()
+            versions["version"] = versions["version"].split("+")[0] + suffix
+            return versions
+
+        versioneer.get_versions = get_versions
+
+
+REQUIRES = ["jax>=0.4.0", "jaxlib>=0.4.0", "jaxtyping"]
 
 EXTRAS = {
     "dev": [