@@ -247,9 +247,9 @@ def __init__(self, buffers: list[str], encoding: str) -> None:
247247 self .deforders : dict [str , int ] = {}
248248 self .finals : list [str ] = []
249249 self .overloads : dict [str , list [Signature ]] = {}
250- self .typing : str | None = None
251- self .typing_final : str | None = None
252- self .typing_overload : str | None = None
250+ self .typing_mods : set [ str ] = set ()
251+ self .typing_final_names : set [ str ] = set ()
252+ self .typing_overload_names : set [ str ] = set ()
253253 super ().__init__ ()
254254
255255 def get_qualname_for (self , name : str ) -> list [str ] | None :
@@ -295,11 +295,8 @@ def add_variable_annotation(self, name: str, annotation: ast.AST) -> None:
295295 self .annotations [basename , name ] = ast_unparse (annotation )
296296
297297 def is_final (self , decorators : list [ast .expr ]) -> bool :
298- final = []
299- if self .typing :
300- final .append ('%s.final' % self .typing )
301- if self .typing_final :
302- final .append (self .typing_final )
298+ final = {f'{ modname } .final' for modname in self .typing_mods }
299+ final |= self .typing_final_names
303300
304301 for decorator in decorators :
305302 try :
@@ -311,11 +308,8 @@ def is_final(self, decorators: list[ast.expr]) -> bool:
311308 return False
312309
313310 def is_overload (self , decorators : list [ast .expr ]) -> bool :
314- overload = []
315- if self .typing :
316- overload .append ('%s.overload' % self .typing )
317- if self .typing_overload :
318- overload .append (self .typing_overload )
311+ overload = {f'{ modname } .overload' for modname in self .typing_mods }
312+ overload |= self .typing_overload_names
319313
320314 for decorator in decorators :
321315 try :
@@ -348,22 +342,24 @@ def visit_Import(self, node: ast.Import) -> None:
348342 for name in node .names :
349343 self .add_entry (name .asname or name .name )
350344
351- if name .name == 'typing' :
352- self .typing = name .asname or name .name
353- elif name .name == 'typing.final' :
354- self .typing_final = name .asname or name .name
355- elif name .name == 'typing.overload' :
356- self .typing_overload = name .asname or name .name
345+ if name .name in { 'typing' , 'typing_extensions' } :
346+ self .typing_mods . add ( name .asname or name .name )
347+ elif name .name in { 'typing.final' , 'typing_extensions.final' } :
348+ self .typing_final_names . add ( name .asname or name .name )
349+ elif name .name in { 'typing.overload' , 'typing_extensions.overload' } :
350+ self .typing_overload_names . add ( name .asname or name .name )
357351
358352 def visit_ImportFrom (self , node : ast .ImportFrom ) -> None :
359353 """Handles Import node and record the order of definitions."""
360354 for name in node .names :
361355 self .add_entry (name .asname or name .name )
362356
363- if node .module == 'typing' and name .name == 'final' :
364- self .typing_final = name .asname or name .name
365- elif node .module == 'typing' and name .name == 'overload' :
366- self .typing_overload = name .asname or name .name
357+ if node .module not in {'typing' , 'typing_extensions' }:
358+ continue
359+ if name .name == 'final' :
360+ self .typing_final_names .add (name .asname or name .name )
361+ elif name .name == 'overload' :
362+ self .typing_overload_names .add (name .asname or name .name )
367363
368364 def visit_Assign (self , node : ast .Assign ) -> None :
369365 """Handles Assign node and pick up a variable comment."""
0 commit comments