Implemented consider-using-dict-items on comprehension

- Also, added new test cases (thanks @cdce8p)

Author: yushao2

pylint/checkers/refactoring/recommendation_checker.py

@@ -1,14 +1,40 @@
 # Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
 # For details: https://github.com/PyCQA/pylint/blob/master/LICENSE
 
-from typing import cast
+from typing import Optional, Union, cast
 
 import astroid
 
 from pylint import checkers, interfaces
 from pylint.checkers import utils
 
 
+def _check_if_dict_keys_used(
+    node: Union[astroid.For, astroid.Comprehension]
+) -> Optional[str]:
+    if not isinstance(node.iter, astroid.Call):
+        # Is it a dictionary?
+        if not isinstance(node.iter, (astroid.Name, astroid.Attribute)):
+            return None
+        inferred = utils.safe_infer(node.iter)
+        if not isinstance(inferred, (astroid.Dict, astroid.DictComp)):
+            return None
+        iterating_object_name = node.iter.as_string()
+
+    else:
+        # Is it a proper keys call?
+        if (
+            isinstance(node.iter.func, astroid.Name)
+            or node.iter.func.attrname != "keys"
+        ):
+            return None
+        inferred = utils.safe_infer(node.iter.func)
+        if not isinstance(inferred, (astroid.BoundMethod, astroid.Dict)):
+            return None
+        iterating_object_name = node.iter.as_string().partition(".")[0]
+    return iterating_object_name
+
+
 class RecommendationChecker(checkers.BaseChecker):
 
     __implements__ = (interfaces.IAstroidChecker,)
@@ -140,26 +166,9 @@ def _check_consider_using_dict_items(self, node: astroid.For) -> None:
         # that the object which is iterated is used as a subscript in the
         # body of the for.
 
-        if not isinstance(node.iter, astroid.Call):
-            # Is it a dictionary?
-            if not isinstance(node.iter, astroid.Name):
-                return
-            inferred = utils.safe_infer(node.iter)
-            if not isinstance(inferred, (astroid.Dict, astroid.DictComp)):
-                return
-            iterating_object_name = node.iter.as_string()
-
-        else:
-            # Is it a proper keys call?
-            if not (
-                isinstance(node.iter.func, astroid.Attribute)
-                and node.iter.func.attrname != "keys"
-            ):
-                return
-            inferred = utils.safe_infer(node.iter.func)
-            if not isinstance(inferred, (astroid.BoundMethod, astroid.Dict)):
-                return
-            iterating_object_name = node.iter.as_string().partition(".")[0]
+        iterating_object_name = _check_if_dict_keys_used(node)
+        if iterating_object_name is None:
+            return
 
         # Verify that the body of the for loop uses a subscript
         # with the object that was iterated. This uses some heuristics
@@ -168,7 +177,8 @@ def _check_consider_using_dict_items(self, node: astroid.For) -> None:
         for child in node.body:
             for subscript in child.nodes_of_class(astroid.Subscript):
                 subscript = cast(astroid.Subscript, subscript)
-                if not isinstance(subscript.value, astroid.Name):
+
+                if not isinstance(subscript.value, (astroid.Name, astroid.Attribute)):
                     continue
 
                 value = subscript.slice
@@ -177,8 +187,9 @@ def _check_consider_using_dict_items(self, node: astroid.For) -> None:
                 if (
                     not isinstance(value, astroid.Name)
                     or value.name != node.target.name
-                    or iterating_object_name != subscript.value.name
+                    or iterating_object_name != subscript.value.as_string()
                 ):
+
                     continue
                 last_definition_lineno = value.lookup(value.name)[1][-1].lineno
                 if last_definition_lineno > node.lineno:
@@ -187,5 +198,39 @@ def _check_consider_using_dict_items(self, node: astroid.For) -> None:
                     # to get the line number where the iterating object was last
                     # defined and compare that to the for loop's line number
                     continue
+                if (
+                    isinstance(subscript.parent, astroid.Assign)
+                    and subscript in subscript.parent.targets
+                ):
+                    # Ignore this subscript if it is the target of an assignment
+                    continue
+
+                self.add_message("consider-using-dict-items", node=node)
+                return
+
+    @utils.check_messages("consider-using-dict-items")
+    def visit_comprehension(self, node: astroid.Comprehension) -> None:
+        iterating_object_name = _check_if_dict_keys_used(node)
+        if iterating_object_name is None:
+            return
+
+        children = list(node.parent.get_children())
+        if node.ifs:
+            children += node.ifs
+        for child in children:
+            for subscript in child.nodes_of_class(astroid.Subscript):
+                subscript = cast(astroid.Subscript, subscript)
+                if not isinstance(subscript.value, (astroid.Name, astroid.Attribute)):
+                    continue
+
+                value = subscript.slice
+                if isinstance(value, astroid.Index):
+                    value = value.value
+                if (
+                    not isinstance(value, astroid.Name)
+                    or value.name != node.target.name
+                    or iterating_object_name != subscript.value.as_string()
+                ):
+                    continue
                 self.add_message("consider-using-dict-items", node=node)
                 return

tests/functional/c/consider/consider_using_dict_items.py

@@ -1,6 +1,6 @@
 """Emit a message for iteration through dict keys and subscripting dict with key."""
 
-# pylint: disable=missing-docstring,unsubscriptable-object
+# pylint: disable=missing-docstring,unsubscriptable-object,too-few-public-methods
 
 def bad():
     a_dict = {1:1, 2:2, 3:3}
@@ -28,3 +28,36 @@ def another_good():
         k = 2
         k = 3
         print(out_of_scope_dict[k])
+
+
+b_dict = {}
+for k2 in b_dict:  # Should not emit warning, key access necessary
+    b_dict[k2] = 2
+
+for k3 in b_dict:  # [consider-using-dict-items]
+    val = b_dict[k3]
+
+for k4 in b_dict.keys():  # [consider-iterating-dictionary,consider-using-dict-items]
+    val = b_dict[k4]
+
+class Foo:
+    c_dict = {}
+
+for k5 in Foo.c_dict:  # [consider-using-dict-items]
+    val = Foo.c_dict[k5]
+
+for k6 in b_dict:  # [consider-using-dict-items]
+    val = b_dict[k6]
+    b_dict[k6] = 2
+
+c_dict={}
+
+val = [(k7, b_dict[k7]) for k7 in b_dict]  # [consider-using-dict-items]
+val = any(True for k8 in b_dict if b_dict[k8])  # [consider-using-dict-items]
+val = {k9: b_dict[k9] for k9 in b_dict}  # [consider-using-dict-items]
+val = [(k7, b_dict[k7]) for k7 in Foo.c_dict]
+val = any(True for k8 in Foo.c_dict if b_dict[k8])
+val = [(k7, c_dict[k7]) for k7 in Foo.c_dict]
+val = any(True for k8 in Foo.c_dict if c_dict[k8])
+val = [(k7, Foo.c_dict[k7]) for k7 in Foo.c_dict] # [consider-using-dict-items]
+val = any(True for k8 in Foo.c_dict if Foo.c_dict[k8]) # [consider-using-dict-items]

tests/functional/c/consider/consider_using_dict_items.txt

@@ -1,3 +1,13 @@
 consider-using-dict-items:7:4:bad:Consider iterating with .items()
 consider-using-dict-items:10:4:bad:Consider iterating with .items()
 consider-using-dict-items:22:4:another_bad:Consider iterating with .items()
+consider-using-dict-items:37:0::Consider iterating with .items()
+consider-iterating-dictionary:40:10::Consider iterating the dictionary directly instead of calling .keys()
+consider-using-dict-items:40:0::Consider iterating with .items()
+consider-using-dict-items:46:0::Consider iterating with .items()
+consider-using-dict-items:49:0::Consider iterating with .items()
+consider-using-dict-items:55:0::Consider iterating with .items()
+consider-using-dict-items:56:0::Consider iterating with .items()
+consider-using-dict-items:57:0::Consider iterating with .items()
+consider-using-dict-items:62:0::Consider iterating with .items()
+consider-using-dict-items:63:0::Consider iterating with .items()