9
9
import itertools
10
10
import operator
11
11
import re
12
+ import sys
12
13
import tokenize
13
14
from token import DEDENT , INDENT , NAME , NEWLINE , NUMBER , OP , STRING
14
15
from tokenize import COMMENT , NL
@@ -338,6 +339,47 @@ def get_line(self, lineno: int) -> str:
338
339
"""Returns specified line."""
339
340
return self .buffers [lineno - 1 ]
340
341
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 :]] + self .buffers [node .lineno :]
352
+ )
353
+ parser .parse ()
354
+ if parser .comment and comment_re .match (parser .comment ):
355
+ for varname in varnames :
356
+ self .add_variable_comment (
357
+ varname , comment_re .sub ('\\ 1' , parser .comment )
358
+ )
359
+ self .add_entry (varname )
360
+ return
361
+
362
+ # check comments before assignment
363
+ if indent_re .match (current_line [: node .col_offset ]):
364
+ comment_lines = []
365
+ for i in range (node .lineno - 1 ):
366
+ before_line = self .get_line (node .lineno - 1 - i )
367
+ if comment_re .match (before_line ):
368
+ comment_lines .append (comment_re .sub ('\\ 1' , before_line ))
369
+ else :
370
+ break
371
+
372
+ if comment_lines :
373
+ comment = dedent_docstring ('\n ' .join (reversed (comment_lines )))
374
+ for varname in varnames :
375
+ self .add_variable_comment (varname , comment )
376
+ self .add_entry (varname )
377
+ return
378
+
379
+ # not commented (record deforders only)
380
+ for varname in varnames :
381
+ self .add_entry (varname )
382
+
341
383
def visit (self , node : ast .AST ) -> None :
342
384
"""Updates self.previous to the given node."""
343
385
super ().visit (node )
@@ -385,52 +427,19 @@ def visit_Assign(self, node: ast.Assign) -> None:
385
427
elif hasattr (node , 'type_comment' ) and node .type_comment :
386
428
for varname in varnames :
387
429
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 :]] + self .buffers [node .lineno :]
392
- )
393
- parser .parse ()
394
- if parser .comment and comment_re .match (parser .comment ):
395
- for varname in varnames :
396
- self .add_variable_comment (
397
- varname , comment_re .sub ('\\ 1' , parser .comment )
398
- )
399
- self .add_entry (varname )
400
- return
401
-
402
- # check comments before assignment
403
- if indent_re .match (current_line [: node .col_offset ]):
404
- comment_lines = []
405
- for i in range (node .lineno - 1 ):
406
- before_line = self .get_line (node .lineno - 1 - i )
407
- if comment_re .match (before_line ):
408
- comment_lines .append (comment_re .sub ('\\ 1' , before_line ))
409
- else :
410
- break
411
-
412
- if comment_lines :
413
- comment = dedent_docstring ('\n ' .join (reversed (comment_lines )))
414
- for varname in varnames :
415
- self .add_variable_comment (varname , comment )
416
- self .add_entry (varname )
417
- return
418
-
419
- # not commented (record deforders only)
420
- for varname in varnames :
421
- self .add_entry (varname )
430
+ self .collect_doc_comment (node , varnames , current_line )
422
431
423
432
def visit_AnnAssign (self , node : ast .AnnAssign ) -> None :
424
433
"""Handles AnnAssign node and pick up a variable comment."""
425
434
self .visit_Assign (node ) # type: ignore[arg-type]
426
435
427
436
def visit_Expr (self , node : ast .Expr ) -> None :
428
437
"""Handles Expr node and pick up a comment if string."""
429
- if (
430
- isinstance (self .previous , ast .Assign | ast .AnnAssign )
431
- and isinstance (node .value , ast .Constant )
432
- and isinstance (node .value .value , str )
438
+ if not (
439
+ isinstance (node .value , ast .Constant ) and isinstance (node .value .value , str )
433
440
):
441
+ return
442
+ if isinstance (self .previous , ast .Assign | ast .AnnAssign ):
434
443
try :
435
444
targets = get_assign_targets (self .previous )
436
445
varnames = get_lvar_names (targets [0 ], self .get_self ())
@@ -444,6 +453,13 @@ def visit_Expr(self, node: ast.Expr) -> None:
444
453
self .add_entry (varname )
445
454
except TypeError :
446
455
pass # this assignment is not new definition!
456
+ if (sys .version_info [:2 ] >= (3 , 12 )) and isinstance (
457
+ self .previous , ast .TypeAlias
458
+ ):
459
+ varname = self .previous .name .id
460
+ docstring = node .value .value
461
+ self .add_variable_comment (varname , dedent_docstring (docstring ))
462
+ self .add_entry (varname )
447
463
448
464
def visit_Try (self , node : ast .Try ) -> None :
449
465
"""Handles Try node and processes body and else-clause.
@@ -488,6 +504,16 @@ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
488
504
"""Handles AsyncFunctionDef node and set context."""
489
505
self .visit_FunctionDef (node ) # type: ignore[arg-type]
490
506
507
+ if sys .version_info [:2 ] >= (3 , 12 ):
508
+ def visit_TypeAlias (self , node : ast .TypeAlias ) -> None :
509
+ """Handles TypeAlias node and picks up a variable comment.
510
+
511
+ .. note:: TypeAlias node refers to `type Foo = Bar` (PEP 695) assignment,
512
+ NOT `Foo: TypeAlias = Bar` (PEP 613).
513
+ """
514
+ current_line = self .get_line (node .lineno )
515
+ self .collect_doc_comment (node , [node .name .id ], current_line )
516
+
491
517
492
518
class DefinitionFinder (TokenProcessor ):
493
519
"""Python source code parser to detect location of functions,
0 commit comments