code_style.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361
  1. # Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
  2. # For details: https://github.com/pylint-dev/pylint/blob/main/LICENSE
  3. # Copyright (c) https://github.com/pylint-dev/pylint/blob/main/CONTRIBUTORS.txt
  4. from __future__ import annotations
  5. import difflib
  6. from typing import TYPE_CHECKING, TypeGuard, cast
  7. from astroid import nodes
  8. from pylint.checkers import BaseChecker, utils
  9. from pylint.checkers.utils import only_required_for_messages, safe_infer
  10. from pylint.interfaces import INFERENCE
  11. if TYPE_CHECKING:
  12. from pylint.lint import PyLinter
  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 regard to performance
  23. and therefore preferred => this is the right place
  24. 3. Everything else should go into another extension
  25. """
  26. name = "code_style"
  27. msgs = {
  28. "R6101": (
  29. "Consider using namedtuple or dataclass for dictionary values",
  30. "consider-using-namedtuple-or-dataclass",
  31. "Emitted when dictionary values can be replaced by namedtuples or dataclass instances.",
  32. ),
  33. "R6102": (
  34. "Consider using an in-place tuple instead of list",
  35. "consider-using-tuple",
  36. "Only for style consistency! "
  37. "Emitted where an in-place defined ``list`` can be replaced by a ``tuple``. "
  38. "Due to optimizations by CPython, there is no performance benefit from it.",
  39. ),
  40. "R6103": (
  41. "Use '%s' instead",
  42. "consider-using-assignment-expr",
  43. "Emitted when an if assignment is directly followed by an if statement and "
  44. "both can be combined by using an assignment expression ``:=``. "
  45. "Requires Python 3.8 and ``py-version >= 3.8``.",
  46. ),
  47. "R6104": (
  48. "Use '%s' to do an augmented assign directly",
  49. "consider-using-augmented-assign",
  50. "Emitted when an assignment is referring to the object that it is assigning "
  51. "to. This can be changed to be an augmented assign.\n"
  52. "Disabled by default!",
  53. {
  54. "default_enabled": False,
  55. },
  56. ),
  57. "R6105": (
  58. "Prefer 'typing.NamedTuple' over 'collections.namedtuple'",
  59. "prefer-typing-namedtuple",
  60. "'typing.NamedTuple' uses the well-known 'class' keyword "
  61. "with type-hints for readability (it's also faster as it avoids "
  62. "an internal exec call).\n"
  63. "Disabled by default!",
  64. {
  65. "default_enabled": False,
  66. },
  67. ),
  68. "R6106": (
  69. "Consider %smath.%s instead of %s",
  70. "consider-math-not-float",
  71. "Using math.inf or math.nan permits to benefit from typing and it is up "
  72. "to 4 times faster than a float call (after the initial import of math). "
  73. "This check also catches typos in float calls as a side effect.",
  74. ),
  75. }
  76. options = (
  77. (
  78. "max-line-length-suggestions",
  79. {
  80. "type": "int",
  81. "default": 0,
  82. "metavar": "<int>",
  83. "help": (
  84. "Max line length for which to sill emit suggestions. "
  85. "Used to prevent optional suggestions which would get split "
  86. "by a code formatter (e.g., black). "
  87. "Will default to the setting for ``max-line-length``."
  88. ),
  89. },
  90. ),
  91. )
  92. def open(self) -> None:
  93. py_version = self.linter.config.py_version
  94. self._py36_plus = py_version >= (3, 6)
  95. self._py38_plus = py_version >= (3, 8)
  96. self._max_length: int = (
  97. self.linter.config.max_line_length_suggestions
  98. or self.linter.config.max_line_length
  99. )
  100. @only_required_for_messages("prefer-typing-namedtuple", "consider-math-not-float")
  101. def visit_call(self, node: nodes.Call) -> None:
  102. if self._py36_plus:
  103. called = safe_infer(node.func)
  104. if not (called and isinstance(called, (nodes.FunctionDef, nodes.ClassDef))):
  105. return
  106. if called.qname() == "collections.namedtuple":
  107. self.add_message(
  108. "prefer-typing-namedtuple", node=node, confidence=INFERENCE
  109. )
  110. elif called.qname() == "builtins.float":
  111. if (
  112. node.args
  113. and isinstance(node.args[0], nodes.Const)
  114. and isinstance(node.args[0].value, str)
  115. and any(
  116. c.isalpha() and c.lower() != "e" for c in node.args[0].value
  117. )
  118. ):
  119. value = node.args[0].value.lower()
  120. math_call: str
  121. if "nan" in value:
  122. math_call = "nan"
  123. elif "inf" in value:
  124. math_call = "inf"
  125. else:
  126. math_call = difflib.get_close_matches(
  127. value, ["inf", "nan"], n=1, cutoff=0
  128. )[0]
  129. minus = "-" if math_call == "inf" and value.startswith("-") else ""
  130. self.add_message(
  131. "consider-math-not-float",
  132. node=node,
  133. args=(minus, math_call, node.as_string()),
  134. confidence=INFERENCE,
  135. )
  136. @only_required_for_messages("consider-using-namedtuple-or-dataclass")
  137. def visit_dict(self, node: nodes.Dict) -> None:
  138. self._check_dict_consider_namedtuple_dataclass(node)
  139. @only_required_for_messages("consider-using-tuple")
  140. def visit_for(self, node: nodes.For) -> None:
  141. if isinstance(node.iter, nodes.List):
  142. self.add_message("consider-using-tuple", node=node.iter)
  143. @only_required_for_messages("consider-using-tuple")
  144. def visit_comprehension(self, node: nodes.Comprehension) -> None:
  145. if isinstance(node.iter, nodes.List):
  146. self.add_message("consider-using-tuple", node=node.iter)
  147. @only_required_for_messages("consider-using-assignment-expr")
  148. def visit_if(self, node: nodes.If) -> None:
  149. if self._py38_plus:
  150. self._check_consider_using_assignment_expr(node)
  151. def _check_dict_consider_namedtuple_dataclass(self, node: nodes.Dict) -> None:
  152. """Check if dictionary values can be replaced by Namedtuple or Dataclass."""
  153. if not (
  154. (
  155. isinstance(node.parent, (nodes.Assign, nodes.AnnAssign))
  156. and isinstance(node.parent.parent, nodes.Module)
  157. )
  158. or (
  159. isinstance(node.parent, nodes.AnnAssign)
  160. and isinstance(node.parent.target, nodes.AssignName)
  161. and utils.is_assign_name_annotated_with(node.parent.target, "Final")
  162. )
  163. ):
  164. # If dict is not part of an 'Assign' or 'AnnAssign' node in
  165. # a module context OR 'AnnAssign' with 'Final' annotation, skip check.
  166. return
  167. # All dict_values are itself dict nodes
  168. if len(node.items) > 1 and all(
  169. isinstance(dict_value, nodes.Dict) for _, dict_value in node.items
  170. ):
  171. KeyTupleT = tuple[type[nodes.NodeNG], str]
  172. # Makes sure all keys are 'Const' string nodes
  173. keys_checked: set[KeyTupleT] = set()
  174. for _, dict_value in node.items:
  175. dict_value = cast(nodes.Dict, dict_value)
  176. for key, _ in dict_value.items:
  177. key_tuple = (type(key), key.as_string())
  178. if key_tuple in keys_checked:
  179. continue
  180. inferred = safe_infer(key)
  181. if not (
  182. isinstance(inferred, nodes.Const)
  183. and inferred.pytype() == "builtins.str"
  184. ):
  185. return
  186. keys_checked.add(key_tuple)
  187. # Makes sure all subdicts have at least 1 common key
  188. key_tuples: list[tuple[KeyTupleT, ...]] = []
  189. for _, dict_value in node.items:
  190. dict_value = cast(nodes.Dict, dict_value)
  191. key_tuples.append(
  192. tuple((type(key), key.as_string()) for key, _ in dict_value.items)
  193. )
  194. keys_intersection: set[KeyTupleT] = set(key_tuples[0])
  195. for sub_key_tuples in key_tuples[1:]:
  196. keys_intersection.intersection_update(sub_key_tuples)
  197. if not keys_intersection:
  198. return
  199. self.add_message("consider-using-namedtuple-or-dataclass", node=node)
  200. return
  201. # All dict_values are itself either list or tuple nodes
  202. if len(node.items) > 1 and all(
  203. isinstance(dict_value, (nodes.List, nodes.Tuple))
  204. for _, dict_value in node.items
  205. ):
  206. # Make sure all sublists have the same length > 0
  207. list_length = len(node.items[0][1].elts)
  208. if list_length == 0:
  209. return
  210. for _, dict_value in node.items[1:]:
  211. if len(dict_value.elts) != list_length:
  212. return
  213. # Make sure at least one list entry isn't a dict
  214. for _, dict_value in node.items:
  215. if all(isinstance(entry, nodes.Dict) for entry in dict_value.elts):
  216. return
  217. self.add_message("consider-using-namedtuple-or-dataclass", node=node)
  218. return
  219. def _check_consider_using_assignment_expr(self, node: nodes.If) -> None:
  220. """Check if an assignment expression (walrus operator) can be used.
  221. For example if an assignment is directly followed by an if statement:
  222. >>> x = 2
  223. >>> if x:
  224. >>> ...
  225. Can be replaced by:
  226. >>> if (x := 2):
  227. >>> ...
  228. Note: Assignment expressions were added in Python 3.8
  229. """
  230. # Check if `node.test` contains a `Name` node
  231. match node.test:
  232. case (
  233. (nodes.Name() as node_name)
  234. | nodes.UnaryOp(op="not", operand=nodes.Name() as node_name)
  235. | nodes.Compare(left=nodes.Name() as node_name, ops=[_])
  236. ):
  237. pass
  238. case _:
  239. return
  240. # Make sure the previous node is an assignment to the same name
  241. # used in `node.test`. Furthermore, ignore if assignment spans multiple lines.
  242. prev_sibling = node.previous_sibling()
  243. if CodeStyleChecker._check_prev_sibling_to_if_stmt(
  244. prev_sibling, node_name.name
  245. ):
  246. # Check if match statement would be a better fit.
  247. # I.e. multiple ifs that test the same name.
  248. if CodeStyleChecker._check_ignore_assignment_expr_suggestion(
  249. node, node_name.name
  250. ):
  251. return
  252. # Build suggestion string. Check length of suggestion
  253. # does not exceed max-line-length-suggestions
  254. test_str = node.test.as_string().replace(
  255. node_name.name,
  256. f"({node_name.name} := {prev_sibling.value.as_string()})",
  257. 1,
  258. )
  259. suggestion = f"if {test_str}:"
  260. if (
  261. node.col_offset is not None
  262. and len(suggestion) + node.col_offset > self._max_length
  263. ) or len(suggestion) > self._max_length:
  264. return
  265. self.add_message(
  266. "consider-using-assignment-expr",
  267. node=node_name,
  268. args=(suggestion,),
  269. )
  270. @staticmethod
  271. def _check_prev_sibling_to_if_stmt(
  272. prev_sibling: nodes.NodeNG | None, name: str | None
  273. ) -> TypeGuard[nodes.Assign | nodes.AnnAssign]:
  274. """Check if previous sibling is an assignment with the same name.
  275. Ignore statements which span multiple lines.
  276. """
  277. if prev_sibling is None or prev_sibling.tolineno - prev_sibling.fromlineno != 0:
  278. return False
  279. match prev_sibling:
  280. case nodes.Assign(
  281. targets=[nodes.AssignName(name=target_name)]
  282. ) | nodes.AnnAssign(target=nodes.AssignName(name=target_name)):
  283. return target_name == name and prev_sibling.value is not None
  284. return False
  285. @staticmethod
  286. def _check_ignore_assignment_expr_suggestion(
  287. node: nodes.If, name: str | None
  288. ) -> bool:
  289. """Return True if suggestion for assignment expr should be ignored.
  290. E.g., in cases where a match statement would be a better fit
  291. (multiple conditions).
  292. """
  293. if isinstance(node.test, nodes.Compare):
  294. next_if_node: nodes.If | None = None
  295. next_sibling = node.next_sibling()
  296. if len(node.orelse) == 1 and isinstance(node.orelse[0], nodes.If):
  297. # elif block
  298. next_if_node = node.orelse[0]
  299. elif isinstance(next_sibling, nodes.If):
  300. # separate if block
  301. next_if_node = next_sibling
  302. match next_if_node:
  303. case nodes.If(
  304. test=nodes.Compare(left=nodes.Name(name=n)) | nodes.Name(name=n)
  305. ) if (n == name):
  306. return True
  307. return False
  308. @only_required_for_messages("consider-using-augmented-assign")
  309. def visit_assign(self, node: nodes.Assign) -> None:
  310. is_aug, op = utils.is_augmented_assign(node)
  311. if is_aug:
  312. self.add_message(
  313. "consider-using-augmented-assign",
  314. args=f"{op}=",
  315. node=node,
  316. line=node.lineno,
  317. col_offset=node.col_offset,
  318. confidence=INFERENCE,
  319. )
  320. def register(linter: PyLinter) -> None:
  321. linter.register_checker(CodeStyleChecker(linter))