Source code for libcst.codemod.visitors._gather_comments

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import re
from typing import Dict, Pattern, Union

import libcst as cst
from libcst.codemod._context import CodemodContext
from libcst.codemod._visitor import ContextAwareVisitor
from libcst.metadata import PositionProvider


[docs] class GatherCommentsVisitor(ContextAwareVisitor): """ Collects all comments matching a certain regex and their line numbers. This visitor is useful for capturing special-purpose comments, for example ``noqa`` style lint suppression annotations. Standalone comments are assumed to affect the line following them, and inline ones are recorded with the line they are on. After visiting a CST, matching comments are collected in the ``comments`` attribute. """ METADATA_DEPENDENCIES = (PositionProvider,) def __init__(self, context: CodemodContext, comment_regex: str) -> None: super().__init__(context) #: Dictionary of comments found in the CST. Keys are line numbers, #: values are comment nodes. self.comments: Dict[int, cst.Comment] = {} self._comment_matcher: Pattern[str] = re.compile(comment_regex) def visit_EmptyLine(self, node: cst.EmptyLine) -> bool: if node.comment is not None: self.handle_comment(node) return False def visit_TrailingWhitespace(self, node: cst.TrailingWhitespace) -> bool: if node.comment is not None: self.handle_comment(node) return False def handle_comment( self, node: Union[cst.EmptyLine, cst.TrailingWhitespace] ) -> None: comment = node.comment assert comment is not None # ensured by callsites above if not self._comment_matcher.match(comment.value): return line = self.get_metadata(PositionProvider, comment).start.line if isinstance(node, cst.EmptyLine): # Standalone comments refer to the next line line += 1 self.comments[line] = comment