code_style.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307
  1. import sys
  2. from typing import List, Optional, Set, Tuple, Type, Union, cast
  3. from astroid import nodes
  4. from pylint.checkers import BaseChecker, utils
  5. from pylint.checkers.utils import check_messages, safe_infer
  6. from pylint.interfaces import IAstroidChecker
  7. from pylint.lint import PyLinter
  8. from pylint.utils.utils import get_global_option
  9. if sys.version_info >= (3, 10):
  10. from typing import TypeGuard
  11. else:
  12. from typing_extensions import TypeGuard
  13. class CodeStyleChecker(BaseChecker):
  14. """Checkers that can improve code consistency.
  15. As such they don't necessarily provide a performance benefit and
  16. are often times opinionated.
  17. Before adding another checker here, consider this:
  18. 1. Does the checker provide a clear benefit,
  19. i.e. detect a common issue or improve performance
  20. => it should probably be part of the core checker classes
  21. 2. Is it something that would improve code consistency,
  22. maybe because it's slightly better with regards to performance
  23. and therefore preferred => this is the right place
  24. 3. Everything else should go into another extension
  25. """
  26. __implements__ = (IAstroidChecker,)
  27. name = "code_style"
  28. priority = -1
  29. msgs = {
  30. "R6101": (
  31. "Consider using namedtuple or dataclass for dictionary values",
  32. "consider-using-namedtuple-or-dataclass",
  33. "Emitted when dictionary values can be replaced by namedtuples or dataclass instances.",
  34. ),
  35. "R6102": (
  36. "Consider using an in-place tuple instead of list",
  37. "consider-using-tuple",
  38. "Only for style consistency! "
  39. "Emitted where an in-place defined ``list`` can be replaced by a ``tuple``. "
  40. "Due to optimizations by CPython, there is no performance benefit from it.",
  41. ),
  42. "R6103": (
  43. "Use '%s' instead",
  44. "consider-using-assignment-expr",
  45. "Emitted when an if assignment is directly followed by an if statement and "
  46. "both can be combined by using an assignment expression ``:=``. "
  47. "Requires Python 3.8 and ``py-version >= 3.8``.",
  48. ),
  49. }
  50. options = (
  51. (
  52. "max-line-length-suggestions",
  53. {
  54. "type": "int",
  55. "metavar": "<int>",
  56. "help": (
  57. "Max line length for which to sill emit suggestions. "
  58. "Used to prevent optional suggestions which would get split "
  59. "by a code formatter (e.g., black). "
  60. "Will default to the setting for ``max-line-length``."
  61. ),
  62. },
  63. ),
  64. )
  65. def __init__(self, linter: PyLinter) -> None:
  66. """Initialize checker instance."""
  67. super().__init__(linter=linter)
  68. def open(self) -> None:
  69. py_version = get_global_option(self, "py-version")
  70. self._py38_plus = py_version >= (3, 8)
  71. self._max_length: int = (
  72. self.config.max_line_length_suggestions
  73. or get_global_option(self, "max-line-length")
  74. )
  75. @check_messages("consider-using-namedtuple-or-dataclass")
  76. def visit_dict(self, node: nodes.Dict) -> None:
  77. self._check_dict_consider_namedtuple_dataclass(node)
  78. @check_messages("consider-using-tuple")
  79. def visit_for(self, node: nodes.For) -> None:
  80. if isinstance(node.iter, nodes.List):
  81. self.add_message("consider-using-tuple", node=node.iter)
  82. @check_messages("consider-using-tuple")
  83. def visit_comprehension(self, node: nodes.Comprehension) -> None:
  84. if isinstance(node.iter, nodes.List):
  85. self.add_message("consider-using-tuple", node=node.iter)
  86. @check_messages("consider-using-assignment-expr")
  87. def visit_if(self, node: nodes.If) -> None:
  88. if self._py38_plus:
  89. self._check_consider_using_assignment_expr(node)
  90. def _check_dict_consider_namedtuple_dataclass(self, node: nodes.Dict) -> None:
  91. """Check if dictionary values can be replaced by Namedtuple or Dataclass."""
  92. if not (
  93. isinstance(node.parent, (nodes.Assign, nodes.AnnAssign))
  94. and isinstance(node.parent.parent, nodes.Module)
  95. or isinstance(node.parent, nodes.AnnAssign)
  96. and isinstance(node.parent.target, nodes.AssignName)
  97. and utils.is_assign_name_annotated_with(node.parent.target, "Final")
  98. ):
  99. # If dict is not part of an 'Assign' or 'AnnAssign' node in
  100. # a module context OR 'AnnAssign' with 'Final' annotation, skip check.
  101. return
  102. # All dict_values are itself dict nodes
  103. if len(node.items) > 1 and all(
  104. isinstance(dict_value, nodes.Dict) for _, dict_value in node.items
  105. ):
  106. KeyTupleT = Tuple[Type[nodes.NodeNG], str]
  107. # Makes sure all keys are 'Const' string nodes
  108. keys_checked: Set[KeyTupleT] = set()
  109. for _, dict_value in node.items:
  110. dict_value = cast(nodes.Dict, dict_value)
  111. for key, _ in dict_value.items:
  112. key_tuple = (type(key), key.as_string())
  113. if key_tuple in keys_checked:
  114. continue
  115. inferred = safe_infer(key)
  116. if not (
  117. isinstance(inferred, nodes.Const)
  118. and inferred.pytype() == "builtins.str"
  119. ):
  120. return
  121. keys_checked.add(key_tuple)
  122. # Makes sure all subdicts have at least 1 common key
  123. key_tuples: List[Tuple[KeyTupleT, ...]] = []
  124. for _, dict_value in node.items:
  125. dict_value = cast(nodes.Dict, dict_value)
  126. key_tuples.append(
  127. tuple((type(key), key.as_string()) for key, _ in dict_value.items)
  128. )
  129. keys_intersection: Set[KeyTupleT] = set(key_tuples[0])
  130. for sub_key_tuples in key_tuples[1:]:
  131. keys_intersection.intersection_update(sub_key_tuples)
  132. if not keys_intersection:
  133. return
  134. self.add_message("consider-using-namedtuple-or-dataclass", node=node)
  135. return
  136. # All dict_values are itself either list or tuple nodes
  137. if len(node.items) > 1 and all(
  138. isinstance(dict_value, (nodes.List, nodes.Tuple))
  139. for _, dict_value in node.items
  140. ):
  141. # Make sure all sublists have the same length > 0
  142. list_length = len(node.items[0][1].elts)
  143. if list_length == 0:
  144. return
  145. for _, dict_value in node.items[1:]:
  146. dict_value = cast(Union[nodes.List, nodes.Tuple], dict_value)
  147. if len(dict_value.elts) != list_length:
  148. return
  149. # Make sure at least one list entry isn't a dict
  150. for _, dict_value in node.items:
  151. dict_value = cast(Union[nodes.List, nodes.Tuple], dict_value)
  152. if all(isinstance(entry, nodes.Dict) for entry in dict_value.elts):
  153. return
  154. self.add_message("consider-using-namedtuple-or-dataclass", node=node)
  155. return
  156. def _check_consider_using_assignment_expr(self, node: nodes.If) -> None:
  157. """Check if an assignment expression (walrus operator) can be used.
  158. For example if an assignment is directly followed by an if statement:
  159. >>> x = 2
  160. >>> if x:
  161. >>> ...
  162. Can be replaced by:
  163. >>> if (x := 2):
  164. >>> ...
  165. Note: Assignment expressions were added in Python 3.8
  166. """
  167. # Check if `node.test` contains a `Name` node
  168. node_name: Optional[nodes.Name] = None
  169. if isinstance(node.test, nodes.Name):
  170. node_name = node.test
  171. elif (
  172. isinstance(node.test, nodes.UnaryOp)
  173. and node.test.op == "not"
  174. and isinstance(node.test.operand, nodes.Name)
  175. ):
  176. node_name = node.test.operand
  177. elif (
  178. isinstance(node.test, nodes.Compare)
  179. and isinstance(node.test.left, nodes.Name)
  180. and len(node.test.ops) == 1
  181. ):
  182. node_name = node.test.left
  183. else:
  184. return
  185. # Make sure the previous node is an assignment to the same name
  186. # used in `node.test`. Furthermore, ignore if assignment spans multiple lines.
  187. prev_sibling = node.previous_sibling()
  188. if CodeStyleChecker._check_prev_sibling_to_if_stmt(
  189. prev_sibling, node_name.name
  190. ):
  191. # Check if match statement would be a better fit.
  192. # I.e. multiple ifs that test the same name.
  193. if CodeStyleChecker._check_ignore_assignment_expr_suggestion(
  194. node, node_name.name
  195. ):
  196. return
  197. # Build suggestion string. Check length of suggestion
  198. # does not exceed max-line-length-suggestions
  199. test_str = node.test.as_string().replace(
  200. node_name.name,
  201. f"({node_name.name} := {prev_sibling.value.as_string()})",
  202. 1,
  203. )
  204. suggestion = f"if {test_str}:"
  205. if (
  206. node.col_offset is not None
  207. and len(suggestion) + node.col_offset > self._max_length
  208. or len(suggestion) > self._max_length
  209. ):
  210. return
  211. self.add_message(
  212. "consider-using-assignment-expr",
  213. node=node_name,
  214. args=(suggestion,),
  215. )
  216. @staticmethod
  217. def _check_prev_sibling_to_if_stmt(
  218. prev_sibling: Optional[nodes.NodeNG], name: Optional[str]
  219. ) -> TypeGuard[Union[nodes.Assign, nodes.AnnAssign]]:
  220. """Check if previous sibling is an assignment with the same name.
  221. Ignore statements which span multiple lines.
  222. """
  223. if prev_sibling is None or prev_sibling.tolineno - prev_sibling.fromlineno != 0:
  224. return False
  225. if (
  226. isinstance(prev_sibling, nodes.Assign)
  227. and len(prev_sibling.targets) == 1
  228. and isinstance(prev_sibling.targets[0], nodes.AssignName)
  229. and prev_sibling.targets[0].name == name
  230. ):
  231. return True
  232. if (
  233. isinstance(prev_sibling, nodes.AnnAssign)
  234. and isinstance(prev_sibling.target, nodes.AssignName)
  235. and prev_sibling.target.name == name
  236. ):
  237. return True
  238. return False
  239. @staticmethod
  240. def _check_ignore_assignment_expr_suggestion(
  241. node: nodes.If, name: Optional[str]
  242. ) -> bool:
  243. """Return True if suggestion for assignment expr should be ignore.
  244. E.g., in cases where a match statement would be a better fit
  245. (multiple conditions).
  246. """
  247. if isinstance(node.test, nodes.Compare):
  248. next_if_node: Optional[nodes.If] = None
  249. next_sibling = node.next_sibling()
  250. if len(node.orelse) == 1 and isinstance(node.orelse[0], nodes.If):
  251. # elif block
  252. next_if_node = node.orelse[0]
  253. elif isinstance(next_sibling, nodes.If):
  254. # separate if block
  255. next_if_node = next_sibling
  256. if ( # pylint: disable=too-many-boolean-expressions
  257. next_if_node is not None
  258. and (
  259. isinstance(next_if_node.test, nodes.Compare)
  260. and isinstance(next_if_node.test.left, nodes.Name)
  261. and next_if_node.test.left.name == name
  262. or isinstance(next_if_node.test, nodes.Name)
  263. and next_if_node.test.name == name
  264. )
  265. ):
  266. return True
  267. return False
  268. def register(linter: PyLinter) -> None:
  269. linter.register_checker(CodeStyleChecker(linter))