99import itertools
1010import operator
1111import re
12+ import sys
1213import tokenize
1314from token import DEDENT , INDENT , NAME , NEWLINE , NUMBER , OP , STRING
1415from tokenize import COMMENT , NL
@@ -338,6 +339,48 @@ def get_line(self, lineno: int) -> str:
338339 """Returns specified line."""
339340 return self .buffers [lineno - 1 ]
340341
342+ def collect_doc_comment (
343+ self ,
344+ # exists for >= 3.12, irrelevant for runtime
345+ node : ast .Assign | ast .TypeAlias , # type: ignore[name-defined]
346+ varnames : list [str ],
347+ current_line : str ,
348+ ) -> None :
349+ # check comments after assignment
350+ parser = AfterCommentParser ([
351+ current_line [node .col_offset :],
352+ * self .buffers [node .lineno :],
353+ ])
354+ parser .parse ()
355+ if parser .comment and comment_re .match (parser .comment ):
356+ for varname in varnames :
357+ self .add_variable_comment (
358+ varname , comment_re .sub ('\\ 1' , parser .comment )
359+ )
360+ self .add_entry (varname )
361+ return
362+
363+ # check comments before assignment
364+ if indent_re .match (current_line [: node .col_offset ]):
365+ comment_lines = []
366+ for i in range (node .lineno - 1 ):
367+ before_line = self .get_line (node .lineno - 1 - i )
368+ if comment_re .match (before_line ):
369+ comment_lines .append (comment_re .sub ('\\ 1' , before_line ))
370+ else :
371+ break
372+
373+ if comment_lines :
374+ comment = dedent_docstring ('\n ' .join (reversed (comment_lines )))
375+ for varname in varnames :
376+ self .add_variable_comment (varname , comment )
377+ self .add_entry (varname )
378+ return
379+
380+ # not commented (record deforders only)
381+ for varname in varnames :
382+ self .add_entry (varname )
383+
341384 def visit (self , node : ast .AST ) -> None :
342385 """Updates self.previous to the given node."""
343386 super ().visit (node )
@@ -385,53 +428,19 @@ def visit_Assign(self, node: ast.Assign) -> None:
385428 elif hasattr (node , 'type_comment' ) and node .type_comment :
386429 for varname in varnames :
387430 self .add_variable_annotation (varname , node .type_comment ) # type: ignore[arg-type]
388-
389- # check comments after assignment
390- parser = AfterCommentParser ([
391- current_line [node .col_offset :],
392- * self .buffers [node .lineno :],
393- ])
394- parser .parse ()
395- if parser .comment and comment_re .match (parser .comment ):
396- for varname in varnames :
397- self .add_variable_comment (
398- varname , comment_re .sub ('\\ 1' , parser .comment )
399- )
400- self .add_entry (varname )
401- return
402-
403- # check comments before assignment
404- if indent_re .match (current_line [: node .col_offset ]):
405- comment_lines = []
406- for i in range (node .lineno - 1 ):
407- before_line = self .get_line (node .lineno - 1 - i )
408- if comment_re .match (before_line ):
409- comment_lines .append (comment_re .sub ('\\ 1' , before_line ))
410- else :
411- break
412-
413- if comment_lines :
414- comment = dedent_docstring ('\n ' .join (reversed (comment_lines )))
415- for varname in varnames :
416- self .add_variable_comment (varname , comment )
417- self .add_entry (varname )
418- return
419-
420- # not commented (record deforders only)
421- for varname in varnames :
422- self .add_entry (varname )
431+ self .collect_doc_comment (node , varnames , current_line )
423432
424433 def visit_AnnAssign (self , node : ast .AnnAssign ) -> None :
425434 """Handles AnnAssign node and pick up a variable comment."""
426435 self .visit_Assign (node ) # type: ignore[arg-type]
427436
428437 def visit_Expr (self , node : ast .Expr ) -> None :
429438 """Handles Expr node and pick up a comment if string."""
430- if (
431- isinstance (self .previous , ast .Assign | ast .AnnAssign )
432- and isinstance (node .value , ast .Constant )
433- and isinstance (node .value .value , str )
439+ if not (
440+ isinstance (node .value , ast .Constant ) and isinstance (node .value .value , str )
434441 ):
442+ return
443+ if isinstance (self .previous , ast .Assign | ast .AnnAssign ):
435444 try :
436445 targets = get_assign_targets (self .previous )
437446 varnames = get_lvar_names (targets [0 ], self .get_self ())
@@ -445,6 +454,13 @@ def visit_Expr(self, node: ast.Expr) -> None:
445454 self .add_entry (varname )
446455 except TypeError :
447456 pass # this assignment is not new definition!
457+ if (sys .version_info [:2 ] >= (3 , 12 )) and isinstance (
458+ self .previous , ast .TypeAlias
459+ ):
460+ varname = self .previous .name .id
461+ docstring = node .value .value
462+ self .add_variable_comment (varname , dedent_docstring (docstring ))
463+ self .add_entry (varname )
448464
449465 def visit_Try (self , node : ast .Try ) -> None :
450466 """Handles Try node and processes body and else-clause.
@@ -489,6 +505,17 @@ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
489505 """Handles AsyncFunctionDef node and set context."""
490506 self .visit_FunctionDef (node ) # type: ignore[arg-type]
491507
508+ if sys .version_info [:2 ] >= (3 , 12 ):
509+
510+ def visit_TypeAlias (self , node : ast .TypeAlias ) -> None :
511+ """Handles TypeAlias node and picks up a variable comment.
512+
513+ .. note:: TypeAlias node refers to `type Foo = Bar` (PEP 695) assignment,
514+ NOT `Foo: TypeAlias = Bar` (PEP 613).
515+ """
516+ current_line = self .get_line (node .lineno )
517+ self .collect_doc_comment (node , [node .name .id ], current_line )
518+
492519
493520class DefinitionFinder (TokenProcessor ):
494521 """Python source code parser to detect location of functions,
0 commit comments