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
@@ -332,6 +333,48 @@ def get_line(self, lineno: int) -> str:
332333 """Returns specified line."""
333334 return self .buffers [lineno - 1 ]
334335
336+ def collect_doc_comment (
337+ self ,
338+ # exists for >= 3.12, irrelevant for runtime
339+ node : ast .Assign | ast .TypeAlias , # type: ignore[name-defined]
340+ varnames : list [str ],
341+ current_line : str ,
342+ ) -> None :
343+ # check comments after assignment
344+ parser = AfterCommentParser ([
345+ current_line [node .col_offset :],
346+ * self .buffers [node .lineno :],
347+ ])
348+ parser .parse ()
349+ if parser .comment and comment_re .match (parser .comment ):
350+ for varname in varnames :
351+ self .add_variable_comment (
352+ varname , comment_re .sub ('\\ 1' , parser .comment )
353+ )
354+ self .add_entry (varname )
355+ return
356+
357+ # check comments before assignment
358+ if indent_re .match (current_line [: node .col_offset ]):
359+ comment_lines = []
360+ for i in range (node .lineno - 1 ):
361+ before_line = self .get_line (node .lineno - 1 - i )
362+ if comment_re .match (before_line ):
363+ comment_lines .append (comment_re .sub ('\\ 1' , before_line ))
364+ else :
365+ break
366+
367+ if comment_lines :
368+ comment = dedent_docstring ('\n ' .join (reversed (comment_lines )))
369+ for varname in varnames :
370+ self .add_variable_comment (varname , comment )
371+ self .add_entry (varname )
372+ return
373+
374+ # not commented (record deforders only)
375+ for varname in varnames :
376+ self .add_entry (varname )
377+
335378 def visit (self , node : ast .AST ) -> None :
336379 """Updates self.previous to the given node."""
337380 super ().visit (node )
@@ -381,53 +424,19 @@ def visit_Assign(self, node: ast.Assign) -> None:
381424 elif hasattr (node , 'type_comment' ) and node .type_comment :
382425 for varname in varnames :
383426 self .add_variable_annotation (varname , node .type_comment ) # type: ignore[arg-type]
384-
385- # check comments after assignment
386- parser = AfterCommentParser ([
387- current_line [node .col_offset :],
388- * self .buffers [node .lineno :],
389- ])
390- parser .parse ()
391- if parser .comment and comment_re .match (parser .comment ):
392- for varname in varnames :
393- self .add_variable_comment (
394- varname , comment_re .sub ('\\ 1' , parser .comment )
395- )
396- self .add_entry (varname )
397- return
398-
399- # check comments before assignment
400- if indent_re .match (current_line [: node .col_offset ]):
401- comment_lines = []
402- for i in range (node .lineno - 1 ):
403- before_line = self .get_line (node .lineno - 1 - i )
404- if comment_re .match (before_line ):
405- comment_lines .append (comment_re .sub ('\\ 1' , before_line ))
406- else :
407- break
408-
409- if comment_lines :
410- comment = dedent_docstring ('\n ' .join (reversed (comment_lines )))
411- for varname in varnames :
412- self .add_variable_comment (varname , comment )
413- self .add_entry (varname )
414- return
415-
416- # not commented (record deforders only)
417- for varname in varnames :
418- self .add_entry (varname )
427+ self .collect_doc_comment (node , varnames , current_line )
419428
420429 def visit_AnnAssign (self , node : ast .AnnAssign ) -> None :
421430 """Handles AnnAssign node and pick up a variable comment."""
422431 self .visit_Assign (node ) # type: ignore[arg-type]
423432
424433 def visit_Expr (self , node : ast .Expr ) -> None :
425434 """Handles Expr node and pick up a comment if string."""
426- if (
427- isinstance (self .previous , ast .Assign | ast .AnnAssign )
428- and isinstance (node .value , ast .Constant )
429- and isinstance (node .value .value , str )
435+ if not (
436+ isinstance (node .value , ast .Constant ) and isinstance (node .value .value , str )
430437 ):
438+ return
439+ if isinstance (self .previous , ast .Assign | ast .AnnAssign ):
431440 try :
432441 targets = get_assign_targets (self .previous )
433442 varnames = get_lvar_names (targets [0 ], self .get_self ())
@@ -441,6 +450,13 @@ def visit_Expr(self, node: ast.Expr) -> None:
441450 self .add_entry (varname )
442451 except TypeError :
443452 pass # this assignment is not new definition!
453+ if (sys .version_info [:2 ] >= (3 , 12 )) and isinstance (
454+ self .previous , ast .TypeAlias
455+ ):
456+ varname = self .previous .name .id
457+ docstring = node .value .value
458+ self .add_variable_comment (varname , dedent_docstring (docstring ))
459+ self .add_entry (varname )
444460
445461 def visit_Try (self , node : ast .Try ) -> None :
446462 """Handles Try node and processes body and else-clause.
@@ -485,6 +501,17 @@ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
485501 """Handles AsyncFunctionDef node and set context."""
486502 self .visit_FunctionDef (node ) # type: ignore[arg-type]
487503
504+ if sys .version_info [:2 ] >= (3 , 12 ):
505+
506+ def visit_TypeAlias (self , node : ast .TypeAlias ) -> None :
507+ """Handles TypeAlias node and picks up a variable comment.
508+
509+ .. note:: TypeAlias node refers to `type Foo = Bar` (PEP 695) assignment,
510+ NOT `Foo: TypeAlias = Bar` (PEP 613).
511+ """
512+ current_line = self .get_line (node .lineno )
513+ self .collect_doc_comment (node , [node .name .id ], current_line )
514+
488515
489516class DefinitionFinder (TokenProcessor ):
490517 """Python source code parser to detect location of functions,
0 commit comments