Run pyright as part of @tb.typecheck, batch all typechecks

pyright doesn't have a cache, so I didn't want to rerun pyright fully
for every test.  So instead I have us generate one file per test case
as part of test suite setup, typecheck all the files at once, and then
check the messages for each file in the individual tests.

This speeds things up for mypy too quite a bit, because loading all
the caches is actually pretty expensive too.

Fixes #878.

@elprans Two tests didn't work for reasons involving splats in shapes.

Author: msullivan

gel/_testbase.py

@@ -1191,7 +1191,86 @@ def gen_lock_key():
     return os.getpid() * 1000 + _lock_cnt
 
 
-def _typecheck(func, *, xfail=False):
+def _typecheck(cls, func, *, xfail=False):
+    wrapped = inspect.unwrap(func)
+    is_async = inspect.iscoroutinefunction(wrapped)
+
+    def run(self):
+        # We already ran the typechecker on everything, so now inspect
+        # the results for this test.
+
+        mypy_output = cls._mypy_errors.get(func.__name__, '')
+        if 'error:' in mypy_output:
+            source_code = _get_file_code(func)
+            lines = source_code.split("\n")
+            pad_width = max(2, len(str(len(lines))))
+            source_code_numbered = "\n".join(
+                f"{i + 1:0{pad_width}d}: {line}"
+                for i, line in enumerate(lines)
+            )
+
+            raise RuntimeError(
+                f"mypy check failed for {func.__name__} "
+                f"\n\ntest code:\n{source_code_numbered}"
+                f"\n\nmypy stdout:\n{mypy_output}"
+            )
+
+        pyright_output = cls._pyright_errors.get(func.__name__, '')
+        if '- error:' in pyright_output:
+            source_code = _get_file_code(func)
+            lines = source_code.split("\n")
+            pad_width = max(2, len(str(len(lines))))
+            source_code_numbered = "\n".join(
+                f"{i + 1:0{pad_width}d}: {line}"
+                for i, line in enumerate(lines)
+            )
+
+            raise RuntimeError(
+                f"pyright check failed for {func.__name__} "
+                f"\n\ntest code:\n{source_code_numbered}"
+                f"\n\npyright stdout:\n{pyright_output}"
+            )
+
+        types = []
+
+        out = mypy_output.split("\n")
+        for line in out:
+            if m := re.match(r'.*Revealed type is "(?P<name>[^"]+)".*', line):
+                types.append(m.group("name"))
+
+        def reveal_type(_, *, ncalls=[0], types=types):  # noqa: B006
+            ncalls[0] += 1
+            try:
+                return types[ncalls[0] - 1]
+            except IndexError:
+                return None
+
+        if is_async:
+            # `func` is the result of `TestCaseMeta.wrap()`, so we
+            # want to use the coroutine function that was there in
+            # the beginning, so let's use `wrapped`
+            wrapped.__globals__["reveal_type"] = reveal_type
+            return wrapped(self)
+        else:
+            func.__globals__["reveal_type"] = reveal_type
+            return func(self)
+
+    if is_async:
+
+        @functools.wraps(wrapped)
+        async def runner(self, run=run):
+            coro = run(self)
+            await coro
+
+        rewrapped = TestCaseMeta.wrap(runner)
+        return unittest.expectedFailure(rewrapped) if xfail else rewrapped
+
+    else:
+        run = functools.wraps(func)(run)
+        return unittest.expectedFailure(run) if xfail else run
+
+
+def _get_file_code(func):
     wrapped = inspect.unwrap(func)
     is_async = inspect.iscoroutinefunction(wrapped)
 
@@ -1235,14 +1314,35 @@ class TestModel(_testbase.{base_class_name}):
 {textwrap.indent(dedented_body, "    " * 2)}
     """
 
-    def run(self):
-        d = type(self).tmp_model_dir.name
+    return source_code
+
+
+def _typecheck_class(cls, funcs):
+    """Extract all the typecheckable functions from a class and typecheck.
+
+    Run both mypy and pyright, then stash the results where the
+    individual functions will deal with them.
+    """
+
+    contents = [(func.__name__, _get_file_code(func)) for func in funcs]
+    cls._mypy_errors = {}
+    cls._pyright_errors = {}
+
+    orig_setUpClass = cls.setUpClass
+
+    def _setUp(cls):
+        orig_setUpClass()
+
+        d = cls.tmp_model_dir.name
 
-        testfn = pathlib.Path(d) / "test.py"
         inifn = pathlib.Path(d) / "mypy.ini"
+        tdir = pathlib.Path(d) / "tests"
+        os.mkdir(tdir)
 
-        with open(testfn, "wt") as f:
-            f.write(source_code)
+        for name, code in contents:
+            testfn = tdir / (name + ".py")
+            with open(testfn, "wt") as f:
+                f.write(code)
 
         with open(inifn, "wt") as f:
             f.write(
@@ -1272,71 +1372,63 @@ def run(self):
                 inifn,
                 "--cache-dir",
                 str(pathlib.Path(__file__).parent.parent / ".mypy_cache"),
-                testfn,
+                tdir,
             ]
-
             res = subprocess.run(
                 cmd,
                 capture_output=True,
                 check=False,
                 cwd=inifn.parent,
             )
-        finally:
-            inifn.unlink()
-            testfn.unlink()
 
-        if res.returncode != 0:
-            lines = source_code.split("\n")
-            pad_width = max(2, len(str(len(lines))))
-            source_code_numbered = "\n".join(
-                f"{i + 1:0{pad_width}d}: {line}"
-                for i, line in enumerate(lines)
-            )
+            cmd = [
+                sys.executable,
+                "-m",
+                "pyright",
+                tdir,
+            ]
 
-            raise RuntimeError(
-                f"mypy check failed for {func.__name__} "
-                f"\n\ntest code:\n{source_code_numbered}"
-                f"\n\nmypy stdout:\n{res.stdout.decode()}"
-                f"\n\nmypy stderr:\n{res.stderr.decode()}"
+            pyright_res = subprocess.run(
+                cmd,
+                capture_output=True,
+                check=False,
+                cwd=inifn.parent,
             )
+        finally:
+            inifn.unlink()
+            shutil.rmtree(tdir)
 
-        types = []
-
-        out = res.stdout.decode().split("\n")
-        for line in out:
-            if m := re.match(r'.*Revealed type is "(?P<name>[^"]+)".*', line):
-                types.append(m.group("name"))
-
-        def reveal_type(_, *, ncalls=[0], types=types):  # noqa: B006
-            ncalls[0] += 1
-            try:
-                return types[ncalls[0] - 1]
-            except IndexError:
-                return None
-
-        if is_async:
-            # `func` is the result of `TestCaseMeta.wrap()`, so we
-            # want to use the coroutine function that was there in
-            # the beginning, so let's use `wrapped`
-            wrapped.__globals__["reveal_type"] = reveal_type
-            return wrapped(self)
-        else:
-            func.__globals__["reveal_type"] = reveal_type
-            return func(self)
-
-    if is_async:
-
-        @functools.wraps(wrapped)
-        async def runner(self, run=run):
-            coro = run(self)
-            await coro
+        # Parse out mypy errors and assign them to test cases.
+        # mypy lines are all prefixed with file name
+        for line in res.stdout.decode('utf-8').split('\n'):
+            start = 'tests' + os.sep
+            if not (line.startswith(start) and '.py' in line):
+                continue
+            name = line.removeprefix(start).split('.')[0]
+            cls._mypy_errors[name] = (
+                cls._mypy_errors.get(name, '') + line + '\n'
+            )
 
-        rewrapped = TestCaseMeta.wrap(runner)
-        return unittest.expectedFailure(rewrapped) if xfail else rewrapped
+        # Parse out mypy errors and assign them to test cases.
+        # Pyright lines have file name groups started by the name, and
+        # then subsequent lines are indented. Messages can be
+        # multiline.  They have a --outputjson mode that oculd save us
+        # trouble here but would give us some more trouble on the
+        # formatting side so whatever.
+        name = None
+        cur_lines = ''
+        for line in pyright_res.stdout.decode('utf-8').split('\n'):
+            if line.startswith('/'):
+                if name:
+                    cls._pyright_errors[name] = cur_lines
+                cur_lines = ''
+                name = line.split('tests' + os.sep)[1].split('.')[0]
+            else:
+                cur_lines += line + '\n'
+        if name:
+            cls._pyright_errors[name] = cur_lines
 
-    else:
-        run = functools.wraps(func)(run)
-        return unittest.expectedFailure(run) if xfail else run
+    cls.setUpClass = classmethod(_setUp)
 
 
 def typecheck(arg):
@@ -1346,20 +1438,21 @@ def typecheck(arg):
     schemas and the query builder APIs.
     """
     # Please don't add arguments to this decorator, thank you.
-    if isinstance(arg, type):
-        for func in arg.__dict__.values():
-            if not isinstance(func, types.FunctionType):
-                continue
-            if not func.__name__.startswith("test_"):
-                continue
-            new_func = typecheck(func)
-            setattr(arg, func.__name__, new_func)
-        return arg
-    else:
-        assert isinstance(arg, types.FunctionType)
+    assert isinstance(arg, type)
+    all_checked = []
+    for func in arg.__dict__.values():
+        if not isinstance(func, types.FunctionType):
+            continue
+        if not func.__name__.startswith("test_"):
+            continue
         if hasattr(arg, "_typecheck_skipped"):
-            return arg
-        return _typecheck(arg)
+            continue
+        all_checked.append(func)
+        new_func = _typecheck(arg, func)
+        setattr(arg, func.__name__, new_func)
+
+    _typecheck_class(arg, all_checked)
+    return arg
 
 
 def skip_typecheck(arg):

tests/test_model_generator.py

@@ -65,7 +65,7 @@ class MockPointer(typing.NamedTuple):
 
 
 @tb.typecheck
-class TestModelGenerator(tb.ModelTestCase):
+class TestModelGeneratorMain(tb.ModelTestCase):
     SCHEMA = os.path.join(os.path.dirname(__file__), "dbsetup", "orm.gel")
 
     SETUP = os.path.join(os.path.dirname(__file__), "dbsetup", "orm.edgeql")
@@ -1575,7 +1575,8 @@ def test_modelgen_data_unpack_polymorphic(self):
 
         q = default.Named.select(
             "*",
-            *default.UserGroup,
+            # FIXME: pyright fails here
+            *default.UserGroup,  # pyright: ignore
         )
 
         for item in self.client.query(q):
@@ -1823,7 +1824,8 @@ def test_modelgen_save_01(self):
 
         pq = (
             default.Post.select(
-                *default.Post,
+                # FIXME: pyright fails here
+                *default.Post,  # pyright: ignore
                 author=True,
             )
             .filter(lambda p: p.body == "I'm Alice")
@@ -4140,9 +4142,12 @@ def test_modelgen_save_reload_links_08(self):
 
         self.assertEqual({u.name for u in g.users}, {"0aaa", "1aaa", "2aaa"})
 
+        u = None
         for u in g.users:
             u.name += "bbb"
 
+        assert u
+
         g.users.remove(u)
         g.users.add(default.User(name="new"))
 
@@ -5088,13 +5093,15 @@ def test_modelgen_linkprops_09(self):
         gs = self.client.get(
             default.GameSession.select("*", players=True).filter(num=123)
         )
+        alice = None
         for p in gs.players:
             if p.name == "Alice":
                 alice = p
                 self.assertEqual(alice.__linkprops__.is_tall_enough, False)
                 alice.__linkprops__.is_tall_enough = True
 
         self.client.sync(gs)
+        assert alice
         self.assertEqual(alice.__linkprops__.is_tall_enough, True)
 
     def test_modelgen_globals_01(self):