From 716031f0e8965288e5c3eee3833db5102f1482ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= <twoertwein@gmail.com>
Date: Fri, 5 Aug 2022 13:02:29 -0400
Subject: [PATCH] stbutest: option to use it with pandas nightly

---
 pyproject.toml           |  4 ++--
 scripts/test/__init__.py |  7 +++++--
 scripts/test/_step.py    |  3 +++
 scripts/test/run.py      | 26 ++++++++++++++++++++++++--
 4 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index e6b72bfcb..07429f6ae 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -89,9 +89,9 @@ help = "Run pyright on 'tests' using the installed stubs"
 script = "scripts.test:test(dist=True, type_checker='pyright')"
 
 [tool.poe.tasks.stubtest]
-script = "scripts.test:stubtest(allowlist, check_missing)"
+script = "scripts.test:stubtest(allowlist, check_missing, nightly)"
 help = "Run stubtest to compare the installed stubs against pandas"
-args = [{ name = "allowlist", positional = true, default = "", required = false, help= "Path to an allowlist (optional)" }, {name = "check_missing", positional = false, default = false, type = "boolean", required = false, help= "Report errors when the stubs are incomplete (off by default)"}]
+args = [{ name = "allowlist", positional = true, default = "", required = false, help= "Path to an allowlist (optional)" }, {name = "check_missing", positional = false, default = false, type = "boolean", required = false, help= "Report errors when the stubs are incomplete (off by default)"}, {name = "nightly", positional = false, default = false, type = "boolean", required = false, help= "Compare against pandas nightly (off by default)"}]
 
 
 [tool.black]
diff --git a/scripts/test/__init__.py b/scripts/test/__init__.py
index 6ba7b65e5..185d7380d 100644
--- a/scripts/test/__init__.py
+++ b/scripts/test/__init__.py
@@ -35,11 +35,14 @@ def test(
     run_job(steps)
 
 
-def stubtest(allowlist: str, check_missing: bool):
+def stubtest(allowlist: str, check_missing: bool, nightly: bool) -> None:
     stubtest = dataclasses.replace(
         _step.stubtest,
         run=partial(
             _step.stubtest.run, allowlist=allowlist, check_missing=check_missing
         ),
     )
-    run_job(_DIST_STEPS[:2] + [stubtest])
+    steps = _DIST_STEPS[:2]
+    if nightly:
+        steps.append(_step.nightly)
+    run_job(steps + [stubtest])
diff --git a/scripts/test/_step.py b/scripts/test/_step.py
index f64fefbbe..56b4d4a02 100644
--- a/scripts/test/_step.py
+++ b/scripts/test/_step.py
@@ -29,3 +29,6 @@
 stubtest = Step(
     name="Run stubtest to compare the installed stubs against pandas", run=run.stubtest
 )
+nightly = Step(
+    name="Install pandas nightly", run=run.nightly_pandas, rollback=run.released_pandas
+)
diff --git a/scripts/test/run.py b/scripts/test/run.py
index 49ec28662..75af694e1 100644
--- a/scripts/test/run.py
+++ b/scripts/test/run.py
@@ -47,7 +47,7 @@ def build_dist():
 
 def install_dist():
     path = sorted(Path("dist/").glob("pandas_stubs-*.whl"))[-1]
-    cmd = ["pip", "install", "--force-reinstall", str(path)]
+    cmd = [sys.executable, "-m", "pip", "install", "--force-reinstall", str(path)]
     subprocess.run(cmd, check=True)
 
 
@@ -69,7 +69,7 @@ def pyright_dist():
 
 
 def uninstall_dist():
-    cmd = ["pip", "uninstall", "-y", "pandas-stubs"]
+    cmd = [sys.executable, "-m", "pip", "uninstall", "-y", "pandas-stubs"]
     subprocess.run(cmd, check=True)
 
 
@@ -78,3 +78,25 @@ def restore_src():
         Path(r"_pandas-stubs").rename("pandas-stubs")
     else:
         raise FileNotFoundError("'_pandas-stubs' folder does not exists.")
+
+
+def nightly_pandas():
+    cmd = [sys.executable, "-m", "pip", "uninstall", "-y", "pandas"]
+    subprocess.run(cmd, check=True)
+    cmd = [
+        sys.executable,
+        "-m",
+        "pip",
+        "install",
+        "-i",
+        "https://pypi.anaconda.org/scipy-wheels-nightly/simple",
+        "pandas",
+    ]
+    subprocess.run(cmd, check=True)
+
+
+def released_pandas():
+    cmd = [sys.executable, "-m", "pip", "uninstall", "-y", "pandas"]
+    subprocess.run(cmd, check=True)
+    cmd = [sys.executable, "-m", "pip", "install", "pandas"]
+    subprocess.run(cmd, check=True)