Check for PEP 604 usage in CI

Since this is a common review issue and our stubs have all been
converted

Author: hauntsaninja

.github/workflows/tests.yml

@@ -17,6 +17,14 @@ jobs:
       - run: pip install toml
       - run: ./tests/check_consistent.py
 
+  pep-604:
+    name: Check for PEP 604 usage
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
+      - run: ./tests/check_pep_604.py
+
   flake8:
     name: Lint with flake8
     runs-on: ubuntu-latest

tests/check_pep_604.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+
+import ast
+import sys
+from pathlib import Path
+
+
+def check_pep_604(tree: ast.AST, path: Path) -> list[str]:
+    errors = []
+
+    class UnionFinder(ast.NodeVisitor):
+        def visit_Subscript(self, node: ast.Subscript) -> None:
+            if (
+                isinstance(node.value, ast.Name)
+                and "Union" in node.value.id.split(".")
+                and isinstance(node.slice, ast.Tuple)
+            ):
+                new_syntax = " | ".join(ast.unparse(x) for x in node.slice.elts)
+                errors.append(
+                    (f"{path}:{node.lineno}: Use PEP 604 syntax for Unions, e.g. `{new_syntax}`")
+                )
+
+    # This doesn't check type aliases (or type var bounds, etc), since those are not
+    # currently supported
+    class AnnotationFinder(ast.NodeVisitor):
+        def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+            UnionFinder().visit(node.annotation)
+
+        def visit_arg(self, node: ast.arg) -> None:
+            if node.annotation is not None:
+                UnionFinder().visit(node.annotation)
+
+    AnnotationFinder().visit(tree)
+    return errors
+
+
+def main() -> None:
+    errors = []
+    for path in Path(".").glob("**/*.pyi"):
+        if "@python2" in path.parts:
+            continue
+        if "stubs/protobuf/google/protobuf" in str(path):  # TODO: fix protobuf stubs
+            continue
+
+        with open(path) as f:
+            tree = ast.parse(f.read())
+        errors.extend(check_pep_604(tree, path))
+
+    if errors:
+        print("\n".join(errors))
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()