B031: allow reusing groups after materialization

README.rst

@@ -500,6 +500,8 @@ Change Log
 UNRELEASED
 ~~~~~~~~~~
 
+* B031: allow reusing a group after assigning ``list(group)`` or ``tuple(group)``
+  back to the same name (#395)
 * B019: also flag `async_lru.alru_cache` and check cache decorators on `async def` methods (#488)
 * B023: don't flag a function whose every reference is a direct call inside the loop body:
   such a function cannot outlive the iteration it was defined in (#468, #380)

bugbear.py

@@ -1362,6 +1362,27 @@ def check_for_b026(self, call: ast.Call) -> None:
             ):
                 self.add_error("B026", starred)
 
+    @staticmethod
+    def _is_b031_group_materialization(node: ast.AST, group_name: str) -> bool:
+        if isinstance(node, ast.Assign):
+            targets = node.targets
+        elif isinstance(node, ast.AnnAssign):
+            targets = [node.target]
+        else:
+            return False
+
+        value = node.value
+        return (
+            any(isinstance(t, ast.Name) and t.id == group_name for t in targets)
+            and isinstance(value, ast.Call)
+            and isinstance(value.func, ast.Name)
+            and value.func.id in {"list", "tuple"}
+            and len(value.args) == 1
+            and isinstance(value.args[0], ast.Name)
+            and value.args[0].id == group_name
+            and not value.keywords
+        )
+
     def _check_b031_group_usages(
         self,
         nodes: Sequence[ast.AST],
@@ -1373,6 +1394,10 @@ def _check_b031_group_usages(
             num_usages = self._check_b031_group_usage(
                 node, group_name, num_usages, repeated
             )
+            if self._is_b031_group_materialization(node, group_name):
+                # The RHS still consumes the generator. After the assignment,
+                # a negative count marks a path where the name is reusable.
+                return -1
         return num_usages
 
     def _check_b031_group_usage(
@@ -1382,6 +1407,9 @@ def _check_b031_group_usage(
         num_usages: int,
         repeated: bool,
     ) -> int:
+        if num_usages < 0:
+            return num_usages
+
         if isinstance(node, ast.Name):
             if node.id == group_name and isinstance(node.ctx, ast.Load):
                 num_usages += 1
@@ -1412,8 +1440,10 @@ def _check_b031_group_usage(
                 node.iter, group_name, num_usages, repeated
             )
             # Any body reference may run once per nested loop iteration.
-            num_usages = self._check_b031_group_usages(
-                node.body, group_name, num_usages, True
+            # A nested loop may never run, so its assignments are not definite.
+            num_usages = max(
+                num_usages,
+                self._check_b031_group_usages(node.body, group_name, num_usages, True),
             )
             return self._check_b031_group_usages(
                 node.orelse, group_name, num_usages, repeated
@@ -1424,16 +1454,20 @@ def _check_b031_group_usage(
                 node.test, group_name, num_usages, repeated
             )
             # A while body can also consume the group on every iteration.
-            num_usages = self._check_b031_group_usages(
-                node.body, group_name, num_usages, True
+            num_usages = max(
+                num_usages,
+                self._check_b031_group_usages(node.body, group_name, num_usages, True),
             )
             return self._check_b031_group_usages(
                 node.orelse, group_name, num_usages, repeated
             )
 
         for child in ast.iter_child_nodes(node):
-            num_usages = self._check_b031_group_usage(
-                child, group_name, num_usages, repeated
+            # Other constructs may contain optional or deferred execution.
+            # Keep usage counts, but do not assume their assignments ran.
+            num_usages = max(
+                num_usages,
+                self._check_b031_group_usage(child, group_name, num_usages, repeated),
             )
         return num_usages
 

tests/eval_files/b031.py

@@ -130,3 +130,89 @@ async def collect_async_groups():
                 collect_shop_items("Jane", section_items)  # B031: 43, "section_items"
             else:
                 collect_shop_items("Joe", section_items)  # B031: 42, "section_items"
+
+
+# Materializing the generator under its original name makes it reusable (#395)
+for _section, section_items in groupby(items, key=lambda p: p[1]):
+    section_items = list(section_items)
+    collect_shop_items("Jane", section_items)
+    collect_shop_items("Joe", section_items)
+
+for _section, section_items in groupby(items, key=lambda p: p[1]):
+    section_items = tuple(section_items)
+    for shopper in shoppers:
+        collect_shop_items(shopper, section_items)
+
+# Annotated and chained assignments also replace the original binding
+for _, group in groupby(items):
+    group: list = list(group)
+    print(group)
+    print(group)
+
+for _, group in groupby(items):
+    saved = group = list(group)
+    print(group)
+    print(saved)
+
+# The generator is reusable after every branch has materialized it
+for _, group in groupby(items):
+    if shoppers:
+        group = list(group)
+    else:
+        group = tuple(group)
+    print(group)
+    print(group)
+
+# A branch that only consumes the generator still makes later uses unsafe
+for _, group in groupby(items):
+    if shoppers:
+        group = list(group)
+    else:
+        print(group)
+    print(group)  # B031: 10, "group"
+
+# A branch can leave the original generator untouched
+for _, group in groupby(items):
+    if shoppers:
+        group = list(group)
+    print(group)
+    print(group)  # B031: 10, "group"
+
+# Materialization must still warn if the generator has already been used
+for _, group in groupby(items):
+    print(group)
+    group = list(group)  # B031: 17, "group"
+    print(group)
+
+# Saving to another name does not make the original generator reusable
+for _, group in groupby(items):
+    saved = list(group)
+    print(group)  # B031: 10, "group"
+
+# Arbitrary calls, including iter(), may return the original iterator
+for _, group in groupby(items):
+    group = iter(group)
+    print(group)  # B031: 10, "group"
+
+# Nested loops may execute zero times, leaving the original generator intact
+for _, group in groupby(items):
+    for _shopper in shoppers:
+        group = list(group)  # B031: 21, "group"
+    print(group)
+    print(group)  # B031: 10, "group"
+
+for _, group in groupby(items):
+    while shoppers:
+        group = tuple(group)  # B031: 22, "group"
+    print(group)
+    print(group)  # B031: 10, "group"
+
+# Assignments in deferred bodies do not replace the enclosing loop variable
+for _, group in groupby(items):
+    def materialize_later(group):
+        if shoppers:
+            group = list(group)
+        else:
+            group = tuple(group)
+    print(group)
+    print(group)  # B031: 10, "group"