diff --git a/README.rst b/README.rst index 5f62d4f..0a16825 100644 --- a/README.rst +++ b/README.rst @@ -506,7 +506,9 @@ UNRELEASED * B020: don't flag `for self.a in self.b`: rebinding an attribute is not rebinding the base name, so two different attributes of the same object are two bindings (#248) * B018: handle also useless calls such as `isinstance(x, int)` without assigning or using the result -* B031: don't count a store-context reference (e.g. an annotation target like `group: T`) as a use of the `groupby` generator (#465) +* B031: don't count a store-context reference (e.g. an annotation target like `group: T`) as a use of the `groupby` generator, + and don't treat references in mutually exclusive ``if``/``elif``/``else`` branches as multiple uses while preserving + warnings when the conditional can run repeatedly (#465) * B902: don't raise a false positive on a metaclass defined with a dotted base such as `abc.ABCMeta` or `enum.EnumMeta` (#411) 25.11.29 diff --git a/bugbear.py b/bugbear.py index 33880ed..d445418 100644 --- a/bugbear.py +++ b/bugbear.py @@ -374,11 +374,6 @@ def children_in_scope(node: ast.AST) -> Iterator[ast.AST]: yield from children_in_scope(child) -def walk_list(nodes: Sequence[ast.AST]) -> Iterator[ast.AST]: - for node in nodes: - yield from ast.walk(node) - - def _typesafe_issubclass(cls: type, class_or_tuple: type | tuple[type, ...]) -> bool: try: return issubclass(cls, class_or_tuple) @@ -1367,6 +1362,81 @@ def check_for_b026(self, call: ast.Call) -> None: ): self.add_error("B026", starred) + def _check_b031_group_usages( + self, + nodes: Sequence[ast.AST], + group_name: str, + num_usages: int = 0, + repeated: bool = False, + ) -> int: + for node in nodes: + num_usages = self._check_b031_group_usage( + node, group_name, num_usages, repeated + ) + return num_usages + + def _check_b031_group_usage( + self, + node: ast.AST, + group_name: str, + num_usages: int, + repeated: bool, + ) -> int: + if isinstance(node, ast.Name): + if node.id == group_name and isinstance(node.ctx, ast.Load): + num_usages += 1 + if repeated or num_usages > 1: + self.add_error("B031", node, node.id) + return num_usages + + if isinstance(node, ast.If): + num_usages = self._check_b031_group_usage( + node.test, group_name, num_usages, repeated + ) + # Only one branch can execute, so keep the largest path count + # instead of adding usages from mutually exclusive branches. + return max( + self._check_b031_group_usages( + node.body, group_name, num_usages, repeated + ), + self._check_b031_group_usages( + node.orelse, group_name, num_usages, repeated + ), + ) + + if isinstance(node, (ast.For, ast.AsyncFor)): + num_usages = self._check_b031_group_usage( + node.target, group_name, num_usages, repeated + ) + num_usages = self._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 + ) + return self._check_b031_group_usages( + node.orelse, group_name, num_usages, repeated + ) + + if isinstance(node, ast.While): + num_usages = self._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 + ) + 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 + ) + return num_usages + def check_for_b031(self, loop_node: ast.For) -> None: # noqa: C901 """Check that `itertools.groupby` isn't iterated over more than once. @@ -1391,30 +1461,7 @@ def check_for_b031(self, loop_node: ast.For) -> None: # noqa: C901 # Ignore any `groupby()` invocation that isn't unpacked return - num_usages = 0 - for node in walk_list(loop_node.body): # type: ignore[assignment] - # Handled nested loops - if isinstance(node, ast.For): - for nested_node in walk_list(node.body): - assert nested_node != node - if ( - isinstance(nested_node, ast.Name) - and nested_node.id == group_name - and isinstance(nested_node.ctx, ast.Load) - ): - self.add_error("B031", nested_node, nested_node.id) - - # Handle multiple uses. Count only loads: a store-context - # reference, such as an annotation target (`group: T`), is - # not a read of the generator (#465). - if ( - isinstance(node, ast.Name) - and node.id == group_name - and isinstance(node.ctx, ast.Load) - ): - num_usages += 1 - if num_usages > 1: - self.add_error("B031", node, node.id) + self._check_b031_group_usages(loop_node.body, group_name) def _get_names_from_tuple(self, node: ast.Tuple) -> Iterator[str]: for dim in node.elts: diff --git a/tests/eval_files/b031.py b/tests/eval_files/b031.py index fd51a70..3c47e8b 100644 --- a/tests/eval_files/b031.py +++ b/tests/eval_files/b031.py @@ -69,3 +69,64 @@ def collect_shop_items(shopper, items): for _section, section_items in groupby(items, key=lambda p: p[1]): section_items: list collect_shop_items("Jane", section_items) + + +# Mutually exclusive branches cannot consume the group more than once (#465) +for _section, section_items in groupby(items, key=lambda p: p[1]): + if _section == "greens": + collect_shop_items("Jane", section_items) + else: + collect_shop_items("Joe", section_items) + +# Each arm of an if/elif/else chain is also mutually exclusive +for _section, section_items in groupby(items, key=lambda p: p[1]): + if _section == "greens": + collect_shop_items("Jane", section_items) + elif _section == "meats & fish": + collect_shop_items("Joe", section_items) + else: + collect_shop_items("Sarah", section_items) + +# Repeated uses on the same path must still warn +for _section, section_items in groupby(items, key=lambda p: p[1]): + if _section == "greens": + collect_shop_items("Jane", section_items) + collect_shop_items("Joe", section_items) # B031: 34, "section_items" + else: + collect_shop_items("Sarah", section_items) + +# A use after a conditional can follow a use inside either branch +for _section, section_items in groupby(items, key=lambda p: p[1]): + if _section == "greens": + collect_shop_items("Jane", section_items) + collect_shop_items("Joe", section_items) # B031: 30, "section_items" + +# A use in the condition happens before either branch +for _section, section_items in groupby(items, key=lambda p: p[1]): + if list(section_items): + collect_shop_items("Jane", section_items) # B031: 35, "section_items" + else: + collect_shop_items("Joe", section_items) # B031: 34, "section_items" + + +# Conditional branches in a repeating while body can run on different iterations +for _section, section_items in groupby(items, key=lambda p: p[1]): + while shoppers: + if _section == "greens": + collect_shop_items("Jane", section_items) # B031: 39, "section_items" + else: + collect_shop_items("Joe", section_items) # B031: 38, "section_items" + + +async def async_shoppers(): + yield "Jane" + + +# The same applies to async for bodies +async def collect_async_groups(): + for _section, section_items in groupby(items, key=lambda p: p[1]): + async for shopper in async_shoppers(): + if shopper == "Jane": + collect_shop_items("Jane", section_items) # B031: 43, "section_items" + else: + collect_shop_items("Joe", section_items) # B031: 42, "section_items"