Merge pull request #549 from numpy/testgen-improvements

Author: jorenham

src/numpy-stubs/@test/generated/literal_bool_ops.pyi

@@ -23,6 +23,7 @@
     TypeVar,
     cast,
     final,
+    get_args,
     overload,
 )
 from typing_extensions import override
@@ -62,12 +63,17 @@
     "xor",
 ]
 _UnOpName: TypeAlias = Literal["abs", "neg", "pos", "invert"]
-_OpName: TypeAlias = Literal[_BinOpName, _UnOpName]
+_OpName: TypeAlias = Literal[_UnOpName, _BinOpName]
 
 ###
 
-ROOT_DIR: Final = Path(__file__).parent.parent
-TARGET_DIR: Final = ROOT_DIR / "src" / "numpy-stubs" / "@test" / "generated"
+DIR_ROOT: Final = Path(__file__).parent.parent
+DIR_SRC: Final = DIR_ROOT / "src"
+DIRS_TARGET: Final = {
+    dir_package.stem: dir_package / "@test" / "generated"
+    for dir_package in DIR_SRC.iterdir()
+    if dir_package.is_dir()
+}
 
 TAB: Final = " " * 4
 BR: Final = "\n"
@@ -366,6 +372,7 @@ def _strip_preamble(source: str) -> tuple[str | None, str]:
 
 
 class TestGen(abc.ABC):
+    package: ClassVar[str]
     stdlib_imports: ClassVar[tuple[str, ...]] = ("from typing import assert_type",)
     numpy_imports: ClassVar[tuple[str, ...]] = (f"import numpy as {NP}",)
 
@@ -383,7 +390,7 @@ def __init__(self) -> None:
     @property
     def path(self) -> Path:
         assert self.testname
-        return TARGET_DIR / f"{self.testname}.pyi"
+        return DIRS_TARGET[self.package] / f"{self.testname}.pyi"
 
     def get_names(self) -> Iterable[tuple[str, str]]:
         return ()
@@ -411,7 +418,7 @@ def _generate_section(self, /, *lines: str) -> Generator[str]:
 
     def _generate_preamble(self) -> Generator[str]:
         timestamp = f"{np.datetime64('now')}Z"
-        here = Path(__file__).relative_to(ROOT_DIR)
+        here = Path(__file__).relative_to(DIR_ROOT)
 
         yield f"# {PREAMBLE_PREFIX} {timestamp} with {here}"
 
@@ -492,7 +499,7 @@ def regenerate(self, /, *, always: bool = False) -> Iterator[str]:
         head_new, body_new = _strip_preamble(src_new)
         assert head_new, src_new
 
-        path_new = str(self.path.relative_to(ROOT_DIR))
+        path_new = str(self.path.relative_to(DIR_ROOT))
         date_new = head_new.split(" ", 1)[0]
 
         if src_old := self._read():
@@ -516,13 +523,14 @@ def regenerate(self, /, *, always: bool = False) -> Iterator[str]:
             tofile=path_new,
             fromfiledate=date_old,
             tofiledate=date_new if write else date_old,
-            n=0,
+            n=1,
             lineterm=BR,
         )
 
 
 @final
 class EMath(TestGen):
+    package = "numpy-stubs"
     testname = "emath"
 
     VALUES: Final[dict[str, list[Any]]] = {
@@ -735,6 +743,7 @@ def get_testcases(self) -> Iterable[str | None]:
 
 @final
 class LiteralBoolOps(TestGen):
+    package = "numpy-stubs"
     testname = "literal_bool_ops"
 
     UNOPS: ClassVar = {
@@ -887,6 +896,7 @@ def get_testcases(self) -> Iterable[str | None]:
 
 @final
 class ScalarOps(TestGen):
+    package = "numpy-stubs"
     testname = "scalar_ops_{}"
 
     OPS_ARITHMETIC: ClassVar[dict[str, _BinOp]] = {
@@ -1144,6 +1154,7 @@ def get_testcases(self) -> Iterable[str | None]:
 
 
 class NDArrayOps(TestGen):
+    package = "numpy-stubs"
     testname = "ndarray_{}"
     numpy_imports_extra: tuple[str, ...] = ("import _numtype as _nt",)
 
@@ -1527,36 +1538,52 @@ def get_testcases(self) -> Iterable[str | None]:
 TESTGENS: Final[Sequence[TestGen]] = [
     EMath(binary=False),
     LiteralBoolOps(),
-    ScalarOps("arithmetic"),
-    ScalarOps("modular"),
-    ScalarOps("bitwise"),
-    ScalarOps("comparison"),
-    NDArrayOps("pos"),
-    NDArrayOps("neg"),
-    NDArrayOps("abs"),
-    NDArrayOps("invert"),
-    NDArrayOps("add"),
-    NDArrayOps("sub"),
-    NDArrayOps("mul"),
-    NDArrayOps("matmul"),
-    NDArrayOps("pow"),
-    NDArrayOps("truediv"),
-    NDArrayOps("floordiv"),
-    NDArrayOps("mod"),
-    NDArrayOps("divmod"),
-    NDArrayOps("lshift"),
-    NDArrayOps("rshift"),
-    NDArrayOps("and"),
-    NDArrayOps("xor"),
-    NDArrayOps("or"),
+    *(ScalarOps(op_kind) for op_kind in get_args(_BinOpKind)),
+    *(NDArrayOps(op_name) for op_name in get_args(_OpName)),
 ]
 
 
 @np.errstate(all="ignore")
 def main() -> None:
-    """(Re)generate the `src/numpy-stubs/@test/generated/{}.pyi` type-tests."""
+    """(Re)generate the `src/*/@test/generated/{}.pyi` type-tests."""
+    cwd = Path.cwd()
+    paths: dict[str, dict[Path, bool]] = {}
+
     for testgen in TESTGENS:
-        sys.stdout.writelines(testgen.regenerate())
+        path = testgen.path
+        diff = testgen.regenerate()
+        diff_out, diff_check = itertools.tee(diff, 2)
+        sys.stderr.writelines(diff_out)
+        sys.stderr.write("\n")
+        sys.stderr.flush()
+
+        diff_count = sum(1 for _ in diff_check)
+        if not diff_count:
+            sys.stdout.write(f"skipped ./{path.relative_to(cwd)}\n")
+            sys.stdout.flush()
+
+        package_paths = paths.setdefault(testgen.package, {})
+        assert path not in package_paths, path
+        package_paths[path] = bool(diff_count)
+
+    orphans: list[Path] = []
+    for package, testdir in DIRS_TARGET.items():
+        if not testdir.exists():
+            continue
+        assert testdir.is_dir()
+
+        known = paths.get(package, {})
+        for path in testdir.rglob("*.pyi"):
+            assert path.is_file()
+            if path not in known:
+                orphans.append(path)
+
+    for orphan in orphans:
+        assert orphan.is_file()
+        orphan.unlink()
+
+        sys.stderr.write(f"removed ./{orphan.relative_to(cwd)}\n")
+        sys.stderr.flush()
 
 
 if __name__ == "__main__":