From f6e4c9c38f1f97a7956440a2b419c26817973836 Mon Sep 17 00:00:00 2001
From: Shantanu <12621235+hauntsaninja@users.noreply.github.com>
Date: Sat, 28 Aug 2021 11:37:50 -0700
Subject: [PATCH] Check for PEP 604 usage in CI (#5903)

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

Co-authored-by: hauntsaninja <>
---
 .github/workflows/tests.yml |  8 ++++
 tests/check_pep_604.py      | 75 +++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100755 tests/check_pep_604.py

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 69ef378df..1dc0ce26b 100644
--- a/.github/workflows/tests.yml
+++ b/.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
diff --git a/tests/check_pep_604.py b/tests/check_pep_604.py
new file mode 100755
index 000000000..a5390dacd
--- /dev/null
+++ b/tests/check_pep_604.py
@@ -0,0 +1,75 @@
+#!/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 node.value.id == "Union"
+                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 Union, e.g. `{new_syntax}`")
+                )
+            if (
+                isinstance(node.value, ast.Name)
+                and node.value.id == "Optional"
+            ):
+                new_syntax = f"{ast.unparse(node.slice)} | None"
+                errors.append(
+                    (f"{path}:{node.lineno}: Use PEP 604 syntax for Optional, 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)
+
+        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+            if node.returns is not None:
+                UnionFinder().visit(node.returns)
+            self.generic_visit(node)
+
+        def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
+            if node.returns is not None:
+                UnionFinder().visit(node.returns)
+            self.generic_visit(node)
+
+    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
+        if "stubs/dateparser/" in str(path):  # TODO: fix dateparser
+            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()