Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/libcst/codemod/visitors/_gather_comments.py: 40%

1# Copyright (c) Meta Platforms, Inc. and affiliates. 

2# 

3# This source code is licensed under the MIT license found in the 

4# LICENSE file in the root directory of this source tree. 

5 

6import re 

7from typing import Dict, Pattern, Union 

8 

9import libcst as cst 

10from libcst.codemod._context import CodemodContext 

11from libcst.codemod._visitor import ContextAwareVisitor 

12from libcst.metadata import PositionProvider 

13 

14 

15class GatherCommentsVisitor(ContextAwareVisitor): 

16 """ 

17 Collects all comments matching a certain regex and their line numbers. 

18 This visitor is useful for capturing special-purpose comments, for example 

19 ``noqa`` style lint suppression annotations. 

20 

21 Standalone comments are assumed to affect the line following them, and 

22 inline ones are recorded with the line they are on. 

23 

24 After visiting a CST, matching comments are collected in the ``comments`` 

25 attribute. 

26 """ 

27 

28 METADATA_DEPENDENCIES = (PositionProvider,) 

29 

30 def __init__(self, context: CodemodContext, comment_regex: str) -> None: 

31 super().__init__(context) 

32 

33 #: Dictionary of comments found in the CST. Keys are line numbers, 

34 #: values are comment nodes. 

35 self.comments: Dict[int, cst.Comment] = {} 

36 

37 self._comment_matcher: Pattern[str] = re.compile(comment_regex) 

38 

39 def visit_EmptyLine(self, node: cst.EmptyLine) -> bool: 

40 if node.comment is not None: 

41 self.handle_comment(node) 

42 return False 

43 

44 def visit_TrailingWhitespace(self, node: cst.TrailingWhitespace) -> bool: 

45 if node.comment is not None: 

46 self.handle_comment(node) 

47 return False 

48 

49 def handle_comment( 

50 self, node: Union[cst.EmptyLine, cst.TrailingWhitespace] 

51 ) -> None: 

52 comment = node.comment 

53 assert comment is not None # ensured by callsites above 

54 if not self._comment_matcher.match(comment.value): 

55 return 

56 line = self.get_metadata(PositionProvider, comment).start.line 

57 if isinstance(node, cst.EmptyLine): 

58 # Standalone comments refer to the next line 

59 line += 1 

60 self.comments[line] = comment