Use PEP 604 syntax wherever possible

[AlexWaygood]

Not currently possible for some unions involving tuple (see python/mypy#11098), and one or two unions involving type[T] and/or Callable.

Script:

#!/usr/bin/env python3

import ast
import re
import subprocess
import sys
from itertools import chain
from pathlib import Path
from typing import NamedTuple


class DeleteableImport(NamedTuple):
    old: str
    replacement: str


def fix_bad_syntax(path: Path) -> None:
    BAD_IMPORTS = {"Union", "Optional"}
    with open(path) as f:
        stub = f.read()

    lines = stub.splitlines()
    tree = ast.parse(stub)
    bad_imports_from_typing = False
    union_alias = "Union"
    optional_alias = "Optional"

    for node in tree.body:
        if isinstance(node, ast.ImportFrom) and node.module == "typing":
            for cls in node.names:
                if cls.name == "Optional":
                    if cls.asname is not None:
                        optional_alias = cls.asname
                    bad_imports_from_typing = True
                elif cls.name == "Union":
                    if cls.asname is not None:
                        union_alias = cls.asname
                    bad_imports_from_typing = True

    if not bad_imports_from_typing:
        return

    errors_found = True

    class OldSyntaxReplacer(ast.NodeVisitor):
        def visit_Subscript(self, node: ast.Subscript) -> None:
            nonlocal errors_found
            subscribed_thing, lineno = node.value, node.lineno

            if isinstance(subscribed_thing, ast.Name) and lineno == node.end_lineno:
                adjusted_lineno = lineno - 1

                if "tuple[" in lines[adjusted_lineno]:  # Avoid https://github.com/python/mypy/issues/11098
                    return None

                subscribed_thing_name, node_slice = subscribed_thing.id, node.slice

                if subscribed_thing_name == union_alias and isinstance(node_slice, ast.Tuple):
                    replacement = " | ".join(ast.unparse(x) for x in node_slice.elts)
                    errors_found = True
                elif subscribed_thing_name == optional_alias:
                    replacement = f"{ast.unparse(node_slice)} | None"
                    errors_found = True
                else:
                    return self.generic_visit(node)

                existing = ast.unparse(node).replace("'", '"')
                lines[adjusted_lineno] = re.sub(re.escape(existing), replacement, lines[adjusted_lineno])
                return None
            else:
                return self.generic_visit(node)

    while errors_found:
        errors_found = False
        OldSyntaxReplacer().visit(ast.parse("\n".join(lines)))

    new_tree = ast.parse("\n".join(lines))

    class SubscriptScanner(ast.NodeVisitor):
        def visit_Subscript(self, node: ast.Subscript) -> None:
            if isinstance(node.value, ast.Name):
                thingname = node.value.id
                if thingname == optional_alias:
                    BAD_IMPORTS.discard("Optional")
                elif thingname == union_alias:
                    BAD_IMPORTS.discard("Union")
            self.generic_visit(node)

    SubscriptScanner().visit(new_tree)
    imports_to_delete = {}

    if BAD_IMPORTS:
        for node in new_tree.body:
            if isinstance(node, ast.ImportFrom) and node.module == "typing":
                new_import_list = [cls for cls in node.names if cls.name not in BAD_IMPORTS]

                ### DEALING WITH EXISTING IMPORT STATEMENTS ###

                # Scenario (1): Now we don't need *any* imports from typing any more.
                if not new_import_list:
                    imports_to_delete[node.lineno - 1] = DeleteableImport(old=ast.unparse(node), replacement="")

                # Scenario (2): we still need imports from typing; the existing import statement is only one line
                elif node.lineno == node.end_lineno:
                    imports_to_delete[node.lineno - 1] = DeleteableImport(
                        old=ast.unparse(node),
                        replacement=ast.unparse(ast.ImportFrom(module="typing", names=new_import_list, level=0)),
                    )

                # Scenario (3): we still need imports from typing; the existing import statement is multiline.
                else:
                    for cls in node.names:
                        if cls.name in BAD_IMPORTS:
                            imports_to_delete[cls.lineno - 1] = DeleteableImport(
                                old=f"{cls.name}," if cls.asname is None else f"{cls.name} as {cls.asname},", replacement=""
                            )

        for lineno, (old_syntax, new_syntax) in imports_to_delete.items():
            lines[lineno] = lines[lineno].replace(old_syntax, new_syntax)

    with open(path, "w") as f:
        f.write("\n".join(lines) + "\n")


def main() -> None:
    print("STARTING RUN: Will attempt to fix new syntax in the stubs directory...\n\n")
    for path in chain(Path("stdlib").rglob("*.pyi"), Path("stubs").rglob("*.pyi")):
        print(f"Attempting to convert {path} to new syntax.")
        fix_bad_syntax(path)

    print("\n\nSTARTING ISORT...\n\n")
    subprocess.run([sys.executable, "-m", "isort", "."])

    print("\n\nSTARTING BLACK...\n\n")
    subprocess.run([sys.executable, "-m", "black", "."])

    print('\n\nRunning "check_new_syntax.py"...\n\n')
    subprocess.run([sys.executable, "tests/check_new_syntax.py"])

    print("\n\nRunning flake8...\n\n")
    subprocess.run([sys.executable, "-m", "flake8", "stdlib"])
    subprocess.run([sys.executable, "-m", "flake8", "stubs"])


if __name__ == "__main__":
    main()

[AlexWaygood]

@erictraut, looks like there's some pyright bugs when it comes to using None in PEP 604 union-type expressions.

[erictraut]

Looking into it now.

[AlexWaygood]

Looking into it now.

Thank you!!

[erictraut]

@AlexWaygood, you hit two bugs in pyright's handling of PEP 604. Both of them involved the heuristics that pyright uses to detect the difference between a variable assignment and an implied type alias declaration (i.e. one that doesn't involve an explicit TypeAlias annotation).

The first bug involves the use of None at the beginning of a union. There is a simple workaround for this one — move the None to the end (or any other position other than the start).

The second bug involves the use of a TypeVar within an implicit type alias that uses PEP 604 syntax. The workaround for this is to add an explicit TypeAlias annotation.

Both of these will be fixed in the next release of pyright (1.1.230) which I plan to publish later today.

[JelleZijlstra]

Unfortunately we can't use TypeAlias in typeshed yet because pytype doesn't support it (#4913). Thanks as always for the quick response @erictraut!

[AlexWaygood]

Let me second @JelleZijlstra — your response time is always incredible @erictraut; thank you for all the time and energy you put into open source software!

[AlexWaygood]

@nipunn1313, this PR makes some stylistic changes to a few of the protobuf stubs. I know some of those are autogenerated -- is modifying those files going to cause problems? (I'm happy to revert the changes to those files, if so!)

[nipunn1313]

Yep! The autogenerated files are autogenerated, so the stylistic updates aren't going to stay permanently. I think it'll be better to revert them out of this diff.

However, I can make the stylistic changes to mypy-protobuf (file me an issue over there), such that we're able to get the benefits over here. I can send a PR to typeshed once we release a mypy-protobuf that includes the PEP604 syntax.

[AlexWaygood]

Yep! The autogenerated files are autogenerated, so the stylistic updates aren't going to stay permanently. I think it'll be better to revert them out of this diff.

However, I can make the stylistic changes to mypy-protobuf (file me an issue over there), such that we're able to get the benefits over here. I can send a PR to typeshed once we release a mypy-protobuf that includes the PEP604 syntax.

Thanks! I've reverted all proposed changes in the stubs/protobuf subdirectory.

Tbh I wouldn't worry too much about updating mypy-protobuf just yet, since there's still a few situations where we can't use the new-style syntax in typeshed due to mypy bugs. (The bugs specifically involve unions that include type[T], tuple, and/or Callable.) So it might be easier just to wait until those bugs are resolved before updating mypy-protobuf!

[srittau]

No CI for you.

[AlexWaygood]

No CI for you.

...Could you possibly try closing and reopening this PR? Is it something to do with me??

[srittau]

No CI for me either. :( But GitHub seemed laggy just now.

[AlexWaygood]

Okay Github Actions is obviously having a bad day.

[AlexWaygood]

Don't know where the mypy_primer comment has got to, but the CI is finally green, and the uploaded mypy_primer diffs are all empty 🥳

[github-actions]

According to mypy_primer, this change has no effect on the checked open source code. 🤖🎉

[srittau]

Not going to check this in depth, but LGTM.

[jakebailey]

Sorry about breaking pyright-action; you should be able to revert the change to the action's version now, I've published a fix.