diff --git a/README.rst b/README.rst index 5a7ce15..d0eec25 100644 --- a/README.rst +++ b/README.rst @@ -501,6 +501,8 @@ Change Log UNRELEASED ~~~~~~~~~~ +* B031: allow reusing a group after assigning ``list(group)`` or ``tuple(group)`` + back to the same name (#395) * B028: report ``warnings.warn`` calls that pass an explicitly empty ``skip_file_prefixes`` tuple (#510) * B019: also flag `async_lru.alru_cache` and check cache decorators on `async def` methods (#488) diff --git a/bugbear.py b/bugbear.py index 71bd71f..a0e8431 100644 --- a/bugbear.py +++ b/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 diff --git a/tests/eval_files/b031.py b/tests/eval_files/b031.py index 3c47e8b..9697086 100644 --- a/tests/eval_files/b031.py +++ b/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"