[pylint] Recurse into subscript subexpressions when searching for list/dict lookups (PLR1733, PLR1736)

crates/ruff_linter/resources/test/fixtures/pylint/unnecessary_dict_index_lookup.py

@@ -38,3 +38,16 @@ def value_intentionally_unused():
         print(FRUITS[fruit_name])  # OK
         blah = FRUITS[fruit_name]  # OK
         assert FRUITS[fruit_name] == "pear"  # OK
+
+
+def rewrite_client_arrays(value_arrays: dict[str, list[int]]) -> dict[str, list[int]]:
+    """Function from https://github.com/zulip/zulip/blob/3da91e951cd03cfa0b9c67378224e348353f36a6/analytics/views/stats.py#L617C1-L626C25"""
+    mapped_arrays: dict[str, list[int]] = {}
+    for label, array in value_arrays.items():
+        mapped_label = client_label_map(label)
+        if mapped_label in mapped_arrays:
+            for i in range(len(array)):
+                mapped_arrays[mapped_label][i] += value_arrays[label][i]  # PLR1733
+        else:
+            mapped_arrays[mapped_label] = [value_arrays[label][i] for i in range(len(array))]  # PLR1733
+    return mapped_arrays

crates/ruff_linter/resources/test/fixtures/pylint/unnecessary_list_index_lookup.py

@@ -76,3 +76,10 @@ def start():
     # PLR1736
     for index, list_item in enumerate(some_list):
         print(some_list[index])
+
+
+def nested_index_lookup():
+    data = {"a": 1, "b": 2}
+    column_names = ["a", "b"]
+    for index, column_name in enumerate(column_names):
+        _ = data[column_names[index]]  # PLR1736

crates/ruff_linter/src/rules/pylint/helpers.rs

@@ -174,27 +174,25 @@ impl<'a> Visitor<'_> for SequenceIndexVisitor<'a> {
         if self.modified {
             return;
         }
-        match expr {
-            Expr::Subscript(ast::ExprSubscript {
-                value,
-                slice,
-                range,
-                ..
-            }) => {
-                let Expr::Name(ast::ExprName { id, .. }) = value.as_ref() else {
-                    return;
-                };
+        if let Expr::Subscript(ast::ExprSubscript {
+            value,
+            slice,
+            range,
+            ..
+        }) = expr
+        {
+            if let Expr::Name(ast::ExprName { id, .. }) = &**value {
                 if id == self.sequence_name {
-                    let Expr::Name(ast::ExprName { id, .. }) = slice.as_ref() else {
-                        return;
-                    };
-                    if id == self.index_name {
-                        self.accesses.push(*range);
+                    if let Expr::Name(ast::ExprName { id, .. }) = &**slice {
+                        if id == self.index_name {
+                            self.accesses.push(*range);
+                        }
                     }
                 }
             }
-            _ => visitor::walk_expr(self, expr),
         }
+
+        visitor::walk_expr(self, expr);
     }
 }
 

crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR1733_unnecessary_dict_index_lookup.py.snap

@@ -121,4 +121,41 @@ unnecessary_dict_index_lookup.py:11:16: PLR1733 [*] Unnecessary lookup of dictio
 13 13 | 
 14 14 | def dont_fix_these():
 
+unnecessary_dict_index_lookup.py:50:51: PLR1733 [*] Unnecessary lookup of dictionary value by key
+   |
+48 |         if mapped_label in mapped_arrays:
+49 |             for i in range(len(array)):
+50 |                 mapped_arrays[mapped_label][i] += value_arrays[label][i]  # PLR1733
+   |                                                   ^^^^^^^^^^^^^^^^^^^ PLR1733
+51 |         else:
+52 |             mapped_arrays[mapped_label] = [value_arrays[label][i] for i in range(len(array))]  # PLR1733
+   |
+   = help: Use existing variable
+
+ℹ Safe fix
+47 47 |         mapped_label = client_label_map(label)
+48 48 |         if mapped_label in mapped_arrays:
+49 49 |             for i in range(len(array)):
+50    |-                mapped_arrays[mapped_label][i] += value_arrays[label][i]  # PLR1733
+   50 |+                mapped_arrays[mapped_label][i] += array[i]  # PLR1733
+51 51 |         else:
+52 52 |             mapped_arrays[mapped_label] = [value_arrays[label][i] for i in range(len(array))]  # PLR1733
+53 53 |     return mapped_arrays
 
+unnecessary_dict_index_lookup.py:52:44: PLR1733 [*] Unnecessary lookup of dictionary value by key
+   |
+50 |                 mapped_arrays[mapped_label][i] += value_arrays[label][i]  # PLR1733
+51 |         else:
+52 |             mapped_arrays[mapped_label] = [value_arrays[label][i] for i in range(len(array))]  # PLR1733
+   |                                            ^^^^^^^^^^^^^^^^^^^ PLR1733
+53 |     return mapped_arrays
+   |
+   = help: Use existing variable
+
+ℹ Safe fix
+49 49 |             for i in range(len(array)):
+50 50 |                 mapped_arrays[mapped_label][i] += value_arrays[label][i]  # PLR1733
+51 51 |         else:
+52    |-            mapped_arrays[mapped_label] = [value_arrays[label][i] for i in range(len(array))]  # PLR1733
+   52 |+            mapped_arrays[mapped_label] = [array[i] for i in range(len(array))]  # PLR1733
+53 53 |     return mapped_arrays

crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR1736_unnecessary_list_index_lookup.py.snap

@@ -218,3 +218,22 @@ unnecessary_list_index_lookup.py:78:15: PLR1736 [*] List index lookup in `enumer
 77 77 |     for index, list_item in enumerate(some_list):
 78    |-        print(some_list[index])
    78 |+        print(list_item)
+79 79 | 
+80 80 | 
+81 81 | def nested_index_lookup():
+
+unnecessary_list_index_lookup.py:85:18: PLR1736 [*] List index lookup in `enumerate()` loop
+   |
+83 |     column_names = ["a", "b"]
+84 |     for index, column_name in enumerate(column_names):
+85 |         _ = data[column_names[index]]  # PLR1736
+   |                  ^^^^^^^^^^^^^^^^^^^ PLR1736
+   |
+   = help: Use the loop variable directly
+
+ℹ Safe fix
+82 82 |     data = {"a": 1, "b": 2}
+83 83 |     column_names = ["a", "b"]
+84 84 |     for index, column_name in enumerate(column_names):
+85    |-        _ = data[column_names[index]]  # PLR1736
+   85 |+        _ = data[column_name]  # PLR1736