Fix crash when using enumerate in a ternary expression (pylint-dev#…

…7132)
Pierre-Sassoulas · Jul 17, 2022 · 3915922 · 3915922
1 parent 9af506a
commit 3915922
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 3 deletions.
diff --git a/doc/whatsnew/2/2.14/full.rst b/doc/whatsnew/2/2.14/full.rst
@@ -6,6 +6,11 @@ What's New in Pylint 2.14.5?
 Release date: TBA
 
 
+* Fixed a crash in the ``undefined-loop-variable`` check when ``enumerate()`` is used
+  in a ternary expression.
+
+  Closes #7131
+
 * Fixed handling of ``--`` as separator between positional arguments and flags.
 
   Closes #7003

diff --git a/pylint/checkers/variables.py b/pylint/checkers/variables.py
@@ -2233,9 +2233,12 @@ def _loopvar_name(self, node: astroid.Name) -> None:
             if (
                 isinstance(inferred, astroid.Instance)
                 and inferred.qname() == "builtins.enumerate"
-                and assign.iter.args
             ):
-                inferred = next(assign.iter.args[0].infer())
+                likely_call = assign.iter
+                if isinstance(assign.iter, nodes.IfExp):
+                    likely_call = assign.iter.body
+                if isinstance(likely_call, nodes.Call):
+                    inferred = next(likely_call.args[0].infer())
         except astroid.InferenceError:
             self.add_message("undefined-loop-variable", args=node.name, node=node)
         else:

diff --git a/tests/functional/u/undefined/undefined_loop_variable.py b/tests/functional/u/undefined/undefined_loop_variable.py
@@ -148,6 +148,13 @@ def use_enumerate():
     print(i, num)
 
 
+def use_enumerate_in_ternary_expression():
+    """https://github.com/PyCQA/pylint/issues/7131"""
+    for i, num in enumerate(range(3)) if __revision__ else enumerate(range(4)):
+        pass
+    print(i, num)
+
+
 def find_even_number(container):
     """https://github.com/PyCQA/pylint/pull/6923#discussion_r895134495"""
     for something in container:

diff --git a/tests/functional/u/undefined/undefined_loop_variable.txt b/tests/functional/u/undefined/undefined_loop_variable.txt
@@ -1,4 +1,4 @@
 undefined-loop-variable:6:11:6:14:do_stuff:Using possibly undefined loop variable 'var':UNDEFINED
 undefined-loop-variable:25:7:25:11::Using possibly undefined loop variable 'var1':UNDEFINED
 undefined-loop-variable:75:11:75:14:do_stuff_with_redefined_range:Using possibly undefined loop variable 'var':UNDEFINED
-undefined-loop-variable:156:11:156:20:find_even_number:Using possibly undefined loop variable 'something':UNDEFINED
+undefined-loop-variable:163:11:163:20:find_even_number:Using possibly undefined loop variable 'something':UNDEFINED