consider_ternary_expression.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. """Check for if / assign blocks that can be rewritten with if-expressions."""
  2. from astroid import nodes
  3. from pylint.checkers import BaseChecker
  4. from pylint.interfaces import IAstroidChecker
  5. class ConsiderTernaryExpressionChecker(BaseChecker):
  6. __implements__ = (IAstroidChecker,)
  7. name = "consider_ternary_expression"
  8. msgs = {
  9. "W0160": (
  10. "Consider rewriting as a ternary expression",
  11. "consider-ternary-expression",
  12. "Multiple assign statements spread across if/else blocks can be "
  13. "rewritten with a single assignment and ternary expression",
  14. )
  15. }
  16. def visit_if(self, node: nodes.If) -> None:
  17. if isinstance(node.parent, nodes.If):
  18. return
  19. if len(node.body) != 1 or len(node.orelse) != 1:
  20. return
  21. bst = node.body[0]
  22. ost = node.orelse[0]
  23. if not isinstance(bst, nodes.Assign) or not isinstance(ost, nodes.Assign):
  24. return
  25. for (bname, oname) in zip(bst.targets, ost.targets):
  26. if not isinstance(bname, nodes.AssignName) or not isinstance(
  27. oname, nodes.AssignName
  28. ):
  29. return
  30. if bname.name != oname.name:
  31. return
  32. self.add_message("consider-ternary-expression", node=node)
  33. def register(linter):
  34. """Required method to auto register this checker.
  35. :param linter: Main interface object for Pylint plugins
  36. :type linter: Pylint object
  37. """
  38. linter.register_checker(ConsiderTernaryExpressionChecker(linter))