for_any_all.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. """Check for use of for loops that only check for a condition."""
  2. from typing import TYPE_CHECKING
  3. from astroid import nodes
  4. from pylint.checkers import BaseChecker
  5. from pylint.checkers.utils import check_messages, returns_bool
  6. from pylint.interfaces import IAstroidChecker
  7. if TYPE_CHECKING:
  8. from pylint.lint.pylinter import PyLinter
  9. class ConsiderUsingAnyOrAllChecker(BaseChecker):
  10. __implements__ = (IAstroidChecker,)
  11. name = "consider-using-any-or-all"
  12. msgs = {
  13. "C0501": (
  14. "`for` loop could be `%s`",
  15. "consider-using-any-or-all",
  16. "A for loop that checks for a condition and return a bool can be replaced with any or all.",
  17. )
  18. }
  19. @check_messages("consider-using-any-or-all")
  20. def visit_for(self, node: nodes.For) -> None:
  21. if len(node.body) != 1: # Only If node with no Else
  22. return
  23. if not isinstance(node.body[0], nodes.If):
  24. return
  25. if_children = list(node.body[0].get_children())
  26. if not len(if_children) == 2: # The If node has only a comparison and return
  27. return
  28. if not returns_bool(if_children[1]):
  29. return
  30. # Check for terminating boolean return right after the loop
  31. node_after_loop = node.next_sibling()
  32. if returns_bool(node_after_loop):
  33. final_return_bool = node_after_loop.value.value
  34. suggested_string = self._build_suggested_string(node, final_return_bool)
  35. self.add_message(
  36. "consider-using-any-or-all", node=node, args=suggested_string
  37. )
  38. @staticmethod
  39. def _build_suggested_string(node: nodes.For, final_return_bool: bool) -> str:
  40. """When a nodes.For node can be rewritten as an any/all statement, return a suggestion for that statement
  41. final_return_bool is the boolean literal returned after the for loop if all conditions fail
  42. """
  43. loop_var = node.target.as_string()
  44. loop_iter = node.iter.as_string()
  45. test_node = next(node.body[0].get_children())
  46. if isinstance(test_node, nodes.UnaryOp) and test_node.op == "not":
  47. # The condition is negated. Advance the node to the operand and modify the suggestion
  48. test_node = test_node.operand
  49. suggested_function = "all" if final_return_bool else "not all"
  50. else:
  51. suggested_function = "not any" if final_return_bool else "any"
  52. test = test_node.as_string()
  53. return f"{suggested_function}({test} for {loop_var} in {loop_iter})"
  54. def register(linter: "PyLinter") -> None:
  55. """Required method to auto register this checker.
  56. :param linter: Main interface object for Pylint plugins
  57. """
  58. linter.register_checker(ConsiderUsingAnyOrAllChecker(linter))