Add consider-ternary-expression extension

6 files changed, 92 insertions, 2 deletions

diff --git a/ChangeLog b/ChangeLog
index 53c870853..31a5609fa 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -69,7 +69,11 @@ Release date: TBA
 
   Closes #4365
 
-* Added optional extension ``while_used``: Emitted whenever a ``while`` loop is used.
+* Added optional extension ``consider-ternary-expression``: Emitted whenever a variable is assigned in both branches of an if/else block.
+
+  Closes # 4366
+
+* Added optional extension ``while-used``: Emitted whenever a ``while`` loop is used.
 
   Closes # 4367
 
diff --git a/doc/whatsnew/2.10.rst b/doc/whatsnew/2.10.rst
index 501c9156a..becc9be44 100644
--- a/doc/whatsnew/2.10.rst
+++ b/doc/whatsnew/2.10.rst
@@ -24,7 +24,11 @@ New checkers
 
   Closes #4365
 
-* Added optional extension ``while_used``: Emitted whenever a ``while`` loop is used.
+* Added optional extension ``consider-ternary-expression``: Emitted whenever a variable is assigned in both branches of an if/else block.
+
+  Closes # 4366
+
+* Added optional extension ``while-used``: Emitted whenever a ``while`` loop is used.
 
   Closes # 4367
 
diff --git a/pylint/extensions/consider_ternary_expression.py b/pylint/extensions/consider_ternary_expression.py
new file mode 100644
index 000000000..976cd20b9
--- /dev/null
+++ b/pylint/extensions/consider_ternary_expression.py
@@ -0,0 +1,53 @@
+"""Check for if / assign blocks that can be rewritten with if-expressions."""
+
+from astroid import nodes
+
+from pylint.checkers import BaseChecker
+from pylint.interfaces import IAstroidChecker
+
+
+class ConsiderTernaryExpressionChecker(BaseChecker):
+
+    __implements__ = (IAstroidChecker,)
+    name = "consider_ternary_expression"
+    msgs = {
+        "W0160": (
+            "Consider rewriting as a ternary expression",
+            "consider-ternary-expression",
+            "Multiple assign statements spread across if/else blocks can be "
+            "rewritten with a single assingment and ternary expression",
+        )
+    }
+
+    def visit_if(self, node: nodes.If) -> None:
+        if isinstance(node.parent, nodes.If):
+            return
+
+        if len(node.body) != 1 or len(node.orelse) != 1:
+            return
+
+        bst = node.body[0]
+        ost = node.orelse[0]
+
+        if not isinstance(bst, nodes.Assign) or not isinstance(ost, nodes.Assign):
+            return
+
+        for (bname, oname) in zip(bst.targets, ost.targets):
+            if not isinstance(bname, nodes.AssignName) or not isinstance(
+                oname, nodes.AssignName
+            ):
+                return
+
+            if bname.name != oname.name:
+                return
+
+        self.add_message("consider-ternary-expression", node=node)
+
+
+def register(linter):
+    """Required method to auto register this checker.
+
+    :param linter: Main interface object for Pylint plugins
+    :type linter: Pylint object
+    """
+    linter.register_checker(ConsiderTernaryExpressionChecker(linter))
diff --git a/tests/functional/ext/consider_ternary_expression.py b/tests/functional/ext/consider_ternary_expression.py
new file mode 100644
index 000000000..5ed5eaac4
--- /dev/null
+++ b/tests/functional/ext/consider_ternary_expression.py
@@ -0,0 +1,17 @@
+if f():  # [consider-ternary-expression]
+    x = 4
+else:
+    x = 5
+
+if g():
+    y = 3
+elif h():
+    y = 4
+else:
+    y = 5
+
+def a():
+    if i():  # [consider-ternary-expression]
+        z = 4
+    else:
+        z = 5
diff --git a/tests/functional/ext/consider_ternary_expression.rc b/tests/functional/ext/consider_ternary_expression.rc
new file mode 100644
index 000000000..39b1d7975
--- /dev/null
+++ b/tests/functional/ext/consider_ternary_expression.rc
@@ -0,0 +1,10 @@
+[MASTER]
+load-plugins=pylint.extensions.consider_ternary_expression,
+
+[MESSAGES CONTROL]
+disable=
+  invalid-name,
+  undefined-variable,
+  unused-variable,
+  missing-function-docstring,
+  missing-module-docstring,
diff --git a/tests/functional/ext/consider_ternary_expression.txt b/tests/functional/ext/consider_ternary_expression.txt
new file mode 100644
index 000000000..08e51e61d
--- /dev/null
+++ b/tests/functional/ext/consider_ternary_expression.txt
@@ -0,0 +1,2 @@
+consider-ternary-expression:1:0::Consider rewriting as a ternary expression
+consider-ternary-expression:14:4:a:Consider rewriting as a ternary expression