More TypeAlias fixes

[AlexWaygood]

I realised that the script I used for preparing #7630 had several bugs.

Here is the new script.

importastimportreimportsubprocessimportsysfromcollectionsimportdefaultdictfromcontextlibimportcontextmanagerfromdataclassesimportdataclassfromitertoolsimportchainfrompathlibimportPathfromtypesimportModuleTypefromtypingimportIterator, NamedTupleFAILURES= []
@dataclassclassNestingCounter:
"""Class to help the PyiVisitor keep track of internal state"""nesting: int=0@contextmanagerdefenabled(self) ->Iterator[None]:
self.nesting+=1try:
yieldfinally:
self.nesting-=1@propertydefactive(self) ->bool:
"""Determine whether the level of nesting is currently non-zero"""returnbool(self.nesting)
classAnnotationFinder(ast.NodeVisitor):
def__init__(self) ->None:
self.all_names=set[str]()
defvisit_Attribute(self, node: ast.Attribute) ->None:
returndefvisit_Name(self, node: ast.Name) ->None:
self.all_names.add(node.id)
self.generic_visit(node)
deffix_bad_syntax(path: Path) ->None:
withopen(path) asf:
stub=f.read()
lines=stub.splitlines()
tree=ast.parse(stub)
typealias_import_needed=FalseclassStubVisitor(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: defaultdict[str, list[ast.Assign]] =defaultdict(list)
# Set of all names and attributes that are used as annotations in the fileself.all_names_in_annotations: set[str] =set()
self.in_class=NestingCounter()
defvisit_ClassDef(self, node: ast.ClassDef) ->None:
withself.in_class.enabled():
self.generic_visit(node)
defvisit_annotation(self, annotation: ast.expr) ->None:
annotation_finder=AnnotationFinder()
annotation_finder.visit(annotation)
self.all_names_in_annotations|=annotation_finder.all_namesdefvisit_FunctionDef(self, node: ast.FunctionDef) ->None:
returns=node.returnsifnode.returnsisnotNone:
self.visit_annotation(node.returns)
self.generic_visit(node)
visit_AsyncFunctionDef=visit_FunctionDefdefvisit_arg(self, node: ast.arg) ->None:
annotation=node.annotationifannotationisnotNone:
self.visit_annotation(annotation)
self.generic_visit(node)
defvisit_AnnAssign(self, node: ast.AnnAssign) ->None:
self.visit_annotation(node)
self.generic_visit(node)
defvisit_Assign(self, node: ast.Assign) ->None:
nonlocaltypealias_import_neededself.generic_visit(node)
target=node.targets[0]
assertisinstance(target, ast.Name)
target_name=target.idifnottarget_name.startswith("_"):
returnif (self.in_class.activeandtarget_name=="__match_args__") or (
target_name=="__all__"andnotself.in_class.active
):
returnassignment=node.valueifisinstance(assignment, (ast.Ellipsis, ast.Call, ast.Num, ast.Str, ast.Bytes)):
returnifisinstance(assignment, ast.Name):
self.maybe_typealias_assignments[target_name].append(node)
elifisinstance(assignment, ast.Attribute):
ifisinstance(assignment.value, ast.Name):
self.maybe_typealias_assignments[target_name].append(node)
elifnot"# noqa: Y026"inlines[node.lineno-1]:
typealias_import_needed=Truelines[node.lineno-1] =re.sub(f"{target_name} = ", f"{target_name}: TypeAlias = ", lines[node.lineno-1])
defvisit(self, tree: ast.AST) ->None:
nonlocaltypealias_import_neededsuper().visit(tree)
forannotationinself.all_names_in_annotations:
fornodeinself.maybe_typealias_assignments[annotation]:
ifnot"# noqa: Y026"inlines[node.lineno-1]:
typealias_import_needed=Truelines[node.lineno-1] =re.sub(f"{annotation} = ", f"{annotation}: TypeAlias = ", lines[node.lineno-1])
StubVisitor().visit(tree)
ifnottypealias_import_needed:
returntree=ast.parse("\n".join(lines))
typealias_imported=FalseclassTypeAliasImportFinder(ast.NodeVisitor):
defvisit_ImportFrom(self, node: ast.ImportFrom) ->None:
nonlocaltypealias_importedifnode.module!="typing_extensions":
returnifany(cls.name=="TypeAlias"forclsinnode.names):
typealias_imported=TruereturnTypeAliasImportFinder().visit(tree)
ifnottypealias_imported:
lines= ["from typing_extensions import TypeAlias"] +lineswithopen(path, "w") asf:
f.write("\n".join(lines) +"\n")
defmain() ->None:
print("STARTING RUN: Will attempt to fix new syntax in the stubs directory...\n\n")
forpathinchain(Path("stdlib").rglob("*.pyi"), Path("stubs").rglob("*.pyi")):
if"@python2"inpath.parts:
continueprint(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", "."])
ifFAILURES:
print("\n\nFAILED to convert the following files to new syntax:\n")
forpathinFAILURES:
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. 🤖🎉