Use TypeAlias where possible for type aliases

[AlexWaygood]

This PR was accomplished using the following steps:

1. Run the below script on typeshed
2. Revert a few places where PEP 612 and PEP 604 combined cause some mypy false-positive errors.
3. Add a few uses of TypeAlias that the script missed, remove a few where the script was overzealous.
4. Manually review the whole thing.

Script I used for step one:

import ast
import re
import subprocess
import sys
from contextlib import contextmanager
from dataclasses import dataclass
from itertools import chain
from pathlib import Path
from types import ModuleType
from typing import Iterator, NamedTuple


FAILURES = []


@dataclass
class NestingCounter:
    """Class to help the PyiVisitor keep track of internal state"""

    nesting: int = 0

    @contextmanager
    def enabled(self) -> Iterator[None]:
        self.nesting += 1
        try:
            yield
        finally:
            self.nesting -= 1

    @property
    def active(self) -> bool:
        """Determine whether the level of nesting is currently non-zero"""
        return bool(self.nesting)


def fix_bad_syntax(path: Path) -> None:
    with open(path) as f:
        stub = f.read()

    lines = stub.splitlines()
    tree = ast.parse(stub)
    typealias_import_needed = False

    class StubVisitor(ast.NodeVisitor):
        def __init__(self) -> None:
            # Mapping of all assignments in the file that could be type aliases
            # (This excludes assignments to function calls and ellipses, etc.)
            self.maybe_typealias_assignments: dict[str, ast.Assign] = {}
            # Set of all names and attributes that are used as annotations in the file
            self.all_annotations: set[str] = set()
            self.in_class = NestingCounter()

        def visit_ClassDef(self, node: ast.ClassDef) -> None:
            with self.in_class.enabled():
                self.generic_visit(node)

        def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
            self.all_annotations.add(ast.unparse(node.annotation))

        def visit_Assign(self, node: ast.Assign) -> None:
            nonlocal typealias_import_needed
            target = node.targets[0]
            assert isinstance(target, ast.Name)
            target_name = target.id
            if (self.in_class.active and target_name == "__match_args__") or (
                target_name == "__all__" and not self.in_class.active
            ):
                return
            assignment = node.value
            if isinstance(assignment, (ast.Ellipsis, ast.Call, ast.Num, ast.Str, ast.Bytes)):
                return

            if isinstance(assignment, ast.Name):
                self.maybe_typealias_assignments[target_name] = node
            elif isinstance(assignment, ast.Attribute):
                if isinstance(assignment.value, ast.Name):
                    self.maybe_typealias_assignments[target_name] = node
            else:
                typealias_import_needed = True
                lines[node.lineno - 1] = re.sub(f"{target_name} = ", f"{target_name}: TypeAlias = ", lines[node.lineno - 1])

        def run(self, tree: ast.AST) -> None:
            nonlocal typealias_import_needed
            self.visit(tree)
            for annotation in self.all_annotations:
                if annotation in self.maybe_typealias_assignments:
                    node = self.maybe_typealias_assignments[annotation]
                    typealias_import_needed = True
                    lines[node.lineno - 1] = re.sub(f"{annotation} = ", f"{annotation}: TypeAlias = ", lines[node.lineno - 1])

    StubVisitor().visit(tree)

    if not typealias_import_needed:
        return

    tree = ast.parse("\n".join(lines))
    typealias_imported = False

    class TypeAliasImportFinder(ast.NodeVisitor):
        def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
            nonlocal typealias_imported

            if node.module != "typing_extensions":
                return

            if any(cls.name == "TypeAlias" for cls in node.names):
                typealias_imported = True
            return

    TypeAliasImportFinder().visit(tree)

    if not typealias_imported:
        lines = ["from typing_extensions import TypeAlias"] + lines

    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")):
        if "@python2" in path.parts:
            continue
        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", "."])

    if FAILURES:
        print("\n\nFAILED to convert the following files to new syntax:\n")
        for path in FAILURES:
            print(f"- {path}")
    else:
        print("\n\nThere were ZERO failures in converting to new syntax. HOORAY!!\n\n")

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


if __name__ == "__main__":
    main()

[github-actions]

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

[JelleZijlstra]

Looks good, thanks for doing all this!

[JelleZijlstra]

This looks a bit odd but seems like they're indeed intended as aliases.

[JelleZijlstra]

good catch, though maybe we should do from builtins import list as _list instead

[AlexWaygood]

Done in #7634!