Skip to content

Conversation

@AlexWaygood
Copy link
Member

@AlexWaygood AlexWaygood commented Dec 24, 2021

Should be easiest to review commit by commit.

This PR was achieved using the following script:

import ast
import re
from collections import defaultdict
from itertools import chain
from operator import attrgetter
from pathlib import Path
from typing import NamedTuple


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


FORBIDDEN_BUILTIN_TYPING_IMPORTS = frozenset({"List", "FrozenSet", "Set", "Dict", "Tuple"})

# AbstractSet intentionally omitted from this list -- special-cased
IMPORTED_FROM_COLLECTIONS_ABC_NOT_TYPING = frozenset(
    {
        "ByteString",
        "Collection",
        "Container",
        "ItemsView",
        "KeysView",
        "Mapping",
        "MappingView",
        "MutableMapping",
        "MutableSequence",
        "MutableSet",
        "Sequence",
        "ValuesView",
        "Iterable",
        "Iterator",
        "Generator",
        "Hashable",
        "Reversible",
        "Sized",
        "Coroutine",
        "AsyncGenerator",
        "AsyncIterable",
        "AsyncIterator",
        "Awaitable",
        "Callable",
    }
)

# The values in the mapping are what these are called in `collections`
IMPORTED_FROM_COLLECTIONS_NOT_TYPING = {
    "Counter": "Counter",
    "Deque": "deque",
    "DefaultDict": "defaultdict",
    "OrderedDict": "OrderedDict",
    "ChainMap": "ChainMap",
}


def fix_bad_syntax(path: Path) -> None:
    if "@python2" in path.parts or Path("stubs/protobuf/google/protobuf") in path.parents:
        return

    with open(path) as f:
        stub = f.read()

    lines = stub.splitlines()
    tree = ast.parse(stub)
    imports_to_delete = {}
    imports_to_add = []
    classes_from_typing = set()
    import_linenos = set()

    class BadImportFinder(ast.NodeVisitor):
        def visit_Import(self, node: ast.Import):
            import_linenos.add(node.lineno)
        
        def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
            import_linenos.add(node.lineno)
            
            if node.module != "typing":
                return

            bad_builtins_classes_in_this_import = set()
            bad_collections_classes_in_this_import = set()
            bad_collections_abc_classes_in_this_import = set()

            for cls in node.names:
                if cls.name in FORBIDDEN_BUILTIN_TYPING_IMPORTS:
                    bad_builtins_classes_in_this_import.add(cls)
                elif cls.name in IMPORTED_FROM_COLLECTIONS_NOT_TYPING:
                    bad_collections_classes_in_this_import.add(cls)
                elif cls.name in IMPORTED_FROM_COLLECTIONS_ABC_NOT_TYPING and path not in {Path("stdlib/_collections_abc.pyi"), Path("stdlib/builtins.pyi")}:
                    bad_collections_abc_classes_in_this_import.add(cls)

            bad_classes_in_this_import = (
                bad_builtins_classes_in_this_import
                | bad_collections_classes_in_this_import
                | bad_collections_abc_classes_in_this_import
            )

            if not bad_classes_in_this_import:
                return

            classes_from_typing.update(cls.name for cls in bad_classes_in_this_import)
            new_import_list = [cls for cls in node.names if cls not in bad_classes_in_this_import]

            if not new_import_list:
                if path == Path("stdlib/csv.pyi"):
                    imports_to_delete[node.lineno - 1] = DeleteableImport(old=ast.unparse(node), replacement="_DictReadMapping = dict")
                elif path != Path("stdlib/collections/__init__.pyi"):
                    imports_to_delete[node.lineno - 1] = DeleteableImport(old=ast.unparse(node), replacement="")
            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)),
                )
            else:
                for cls in node.names:
                    if cls in bad_classes_in_this_import:
                        imports_to_delete[cls.lineno - 1] = DeleteableImport(
                            old=f"{cls.name}," if cls.asname is None else f"{cls.name} as {cls.asname},",
                            replacement=""
                        )

            if bad_collections_classes_in_this_import:
                imports_to_add.append(
                    ast.unparse(
                        ast.ImportFrom(
                            module="collections",
                            names=[
                                ast.alias(name=IMPORTED_FROM_COLLECTIONS_NOT_TYPING[cls.name], asname=cls.asname)
                                for cls in sorted(bad_collections_classes_in_this_import, key=attrgetter("name"))
                            ],
                            level=0,
                        )
                    )
                )

            if bad_collections_abc_classes_in_this_import and path != Path("stdlib/collections/__init__.pyi"):
                imports_to_add.append(
                    ast.unparse(
                        ast.ImportFrom(
                            module="collections.abc",
                            names=sorted(bad_collections_abc_classes_in_this_import, key=attrgetter("name")),
                            level=0,
                        )
                    )
                )

    BadImportFinder().visit(tree)

    if not classes_from_typing:
        return

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

    first_import_lineno = min(import_linenos) - 1
    
    for new_import in imports_to_add:
        lines[first_import_lineno:first_import_lineno] = [new_import]

    try:
        new_tree = ast.parse("\n".join(lines))
    except SyntaxError:
        print(path)
    else:
        lines_with_bad_syntax = defaultdict(list)

        class OldSyntaxFinder(ast.NodeVisitor):
            def visit_Subscript(self, node: ast.Subscript) -> None:
                if isinstance(node.value, ast.Name) and node.value.id in (
                    classes_from_typing & (FORBIDDEN_BUILTIN_TYPING_IMPORTS | {"Deque", "DefaultDict"})
                ):
                    lines_with_bad_syntax[node.lineno - 1].append(node.value.id)
                self.generic_visit(node)

        OldSyntaxFinder().visit(new_tree)

        for i, cls_list in lines_with_bad_syntax.items():
            for cls in cls_list:
                lines[i] = re.sub(fr"(\W){cls}\[", fr"\1{cls.lower()}[", lines[i])

    if Path("stubs\typed-ast\typed_ast") in path.parents:
        lines.remove("import typing")

    new_stub = '\n'.join(lines) + '\n'

    if path == Path("stdlib/plistlib.pyi"):
        new_stub = new_stub.replace("_Dict", "dict")

    with open(path, "w") as f:
        f.write(new_stub)


def main() -> None:
    for path in chain(Path("stdlib").rglob("*.pyi"), Path("stubs").rglob("*.pyi")):
        fix_bad_syntax(path)


if __name__ == "__main__":
    main()

@AlexWaygood AlexWaygood reopened this Dec 24, 2021
@github-actions

This comment has been minimized.

1 similar comment
@github-actions

This comment has been minimized.

@github-actions
Copy link
Contributor

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

@AlexWaygood AlexWaygood marked this pull request as ready for review December 24, 2021 23:27
@AlexWaygood
Copy link
Member Author

Third-party stubtest errors seem unrelated to this PR.

@srittau
Copy link
Collaborator

srittau commented Dec 24, 2021

But the pytype error needs to be fixed. stubtest fix is in #6685.

@AlexWaygood
Copy link
Member Author

But the pytype error needs to be fixed.

Yeah, of course. No rush :)

@AlexWaygood
Copy link
Member Author

There is now a merge conflict in this PR, and I have also spotted a few other errors as well. I will close this for now, and try again soon.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants