B911: don't flag a bare batched() the module binds to something else

README.rst

@@ -615,6 +615,9 @@ UNRELEASED
   branches inside loops (#606)
 * B042: don't treat a class named like an exception as one when its only base is
   ``object``, as in ``class ParsedException(object)`` (#538)
+* B911: only treat a bare ``batched()`` call as ``itertools.batched`` when the
+  module doesn't bind ``batched`` to something else, such as its own
+  ``def batched``, the same way B031 handles ``groupby()``
 
 26.9.30
 ~~~~~~~

bugbear.py

@@ -707,21 +707,43 @@ def _b906_definition_exprs(
     return exprs
 
 
-def _b031_module_groupby_is_itertools(module: ast.Module) -> bool:
-    """Whether a bare ``groupby`` is ``itertools.groupby`` or never bound.
+def _module_assigns_itertools_attr(module: ast.Module, name: str) -> bool:
+    """Whether the module has ``name = itertools.<name>``, e.g. in a fallback."""
+    itertools_modules: set[str] = set()
+    bases: set[str] = set()
+    for node in ast.walk(module):
+        if isinstance(node, ast.Import):
+            itertools_modules.update(
+                alias.asname or alias.name
+                for alias in node.names
+                if alias.name == "itertools"
+            )
+        elif (
+            isinstance(node, ast.Assign)
+            and isinstance(node.value, ast.Attribute)
+            and node.value.attr == name
+            and isinstance(node.value.value, ast.Name)
+            and any(isinstance(t, ast.Name) and t.id == name for t in node.targets)
+        ):
+            bases.add(node.value.value.id)
+    return not itertools_modules.isdisjoint(bases)
+
+
+def _module_name_is_itertools(module: ast.Module, name: str) -> bool:
+    """Whether a bare ``name`` is ``itertools.<name>`` or never bound.
 
     Whole-module ``ast.walk``; nested scopes are intentionally not modeled.
     """
     from_itertools = False
     other = False
     for node in ast.walk(module):
         if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
-            if node.name == "groupby":
+            if node.name == name:
                 other = True
         elif isinstance(node, (ast.Import, ast.ImportFrom)):
             for alias in node.names:
-                # Absolute ``from itertools import *`` binds ``groupby``.
-                # Other star imports do not. ``import groupby.sub`` does.
+                # Absolute ``from itertools import *`` binds ``name``.
+                # Other star imports do not. ``import name.sub`` does.
                 if alias.name == "*":
                     if (
                         isinstance(node, ast.ImportFrom)
@@ -731,29 +753,29 @@ def _b031_module_groupby_is_itertools(module: ast.Module) -> bool:
                         from_itertools = True
                     continue
                 bound = alias.asname or alias.name.split(".", 1)[0]
-                if bound != "groupby":
+                if bound != name:
                     continue
                 if (
                     isinstance(node, ast.ImportFrom)
                     and node.level == 0
                     and node.module == "itertools"
-                    and alias.name == "groupby"
+                    and alias.name == name
                 ):
                     from_itertools = True
                 else:
                     other = True
         elif (
             isinstance(node, ast.Name)
-            and node.id == "groupby"
+            and node.id == name
             and isinstance(node.ctx, ast.Store)
         ):
             other = True
         elif (
             isinstance(node, (ast.MatchAs, ast.MatchStar, ast.ExceptHandler))
-            and node.name == "groupby"
-        ) or (isinstance(node, ast.MatchMapping) and node.rest == "groupby"):
+            and node.name == name
+        ) or (isinstance(node, ast.MatchMapping) and node.rest == name):
             other = True
-    return from_itertools or not other
+    return from_itertools or not other or _module_assigns_itertools_attr(module, name)
 
 
 @attr.s
@@ -783,8 +805,8 @@ class BugBearVisitor(ast.NodeVisitor):
     _b008_imports: dict[str, str | None] = attr.ib(factory=dict, init=False)
     _b008_class_imports: list[dict[str, str | None]] = attr.ib(factory=list, init=False)
     _b008_class_globals: list[set[str]] = attr.ib(factory=list, init=False)
-    # None until the enclosing module has been scanned for bindings of groupby.
-    _b031_bare_groupby_is_itertools: bool | None = attr.ib(default=None, init=False)
+    # Bare itertools names (groupby, batched) already resolved for this module.
+    _bare_itertools_names: dict[str, bool] = attr.ib(factory=dict, init=False)
 
     # set to "*" when inside a try/except*, for correctly printing errors
     in_trystar: str = attr.ib(default="")
@@ -2146,17 +2168,17 @@ def _check_b031_group_usage(
             )
         return num_usages
 
-    def _b031_groupby_is_itertools(self) -> bool:
-        """Whether a bare ``groupby()`` in this module is ``itertools.groupby``."""
-        if self._b031_bare_groupby_is_itertools is None:
+    def _bare_name_is_itertools(self, name: str) -> bool:
+        """Whether a bare ``name()`` in this module is ``itertools.<name>``."""
+        if name not in self._bare_itertools_names:
             module = self.contexts[0].node if self.contexts else None
             # No module to inspect: keep the historical report.
-            self._b031_bare_groupby_is_itertools = (
-                _b031_module_groupby_is_itertools(module)
+            self._bare_itertools_names[name] = (
+                _module_name_is_itertools(module, name)
                 if isinstance(module, ast.Module)
                 else True
             )
-        return self._b031_bare_groupby_is_itertools
+        return self._bare_itertools_names[name]
 
     def check_for_b031(self, loop_node: ast.For) -> None:  # noqa: C901
         """Check that `itertools.groupby` isn't iterated over more than once.
@@ -2171,7 +2193,7 @@ def check_for_b031(self, loop_node: ast.For) -> None:  # noqa: C901
             if (
                 isinstance(node.func, ast.Name)
                 and node.func.id == "groupby"
-                and self._b031_groupby_is_itertools()
+                and self._bare_name_is_itertools("groupby")
             ) or (
                 isinstance(node.func, ast.Attribute)
                 and node.func.attr == "groupby"
@@ -3219,7 +3241,11 @@ def check_for_b910(self, node: ast.Call) -> None:
 
     def check_for_b911(self, node: ast.Call) -> None:
         if (
-            (isinstance(node.func, ast.Name) and node.func.id == "batched")
+            (
+                isinstance(node.func, ast.Name)
+                and node.func.id == "batched"
+                and self._bare_name_is_itertools("batched")
+            )
             or (
                 isinstance(node.func, ast.Attribute)
                 and node.func.attr == "batched"

tests/eval_files/b031_shadowed_alias.py

@@ -0,0 +1,15 @@
+"""
+Should emit:
+B031 - on line 14
+"""
+
+# Assigning itertools.groupby to groupby still means itertools.groupby.
+import itertools
+
+groupby = itertools.groupby
+
+for _, v in groupby(None, None):
+    if v:
+        pass
+    if v:  # B031: 7, "v"
+        pass

tests/eval_files/b911_shadowed.py

@@ -0,0 +1,16 @@
+"""
+Should emit:
+B911 - on line 16
+"""
+
+# A module-level function named batched is not itertools.batched.
+import itertools
+
+
+def batched(iterable, n):
+    it = iter(iterable)
+    return iter(lambda: tuple(itertools.islice(it, n)), ())
+
+
+batched(range(3), 2)
+itertools.batched(range(3), 2)  # B911: 0

tests/eval_files/b911_shadowed_alias.py

@@ -0,0 +1,11 @@
+"""
+Should emit:
+B911 - on line 11
+"""
+
+# Assigning itertools.batched to batched still means itertools.batched.
+import itertools
+
+batched = itertools.batched
+
+batched(range(3), 2)  # B911: 0

tests/eval_files/b911_shadowed_alias_fallback.py

@@ -0,0 +1,18 @@
+"""
+Should emit:
+B911 - on line 18
+"""
+
+# A fallback for older Pythons still means batched can be itertools.batched.
+import itertools
+
+try:
+    batched = itertools.batched
+except AttributeError:
+
+    def batched(iterable, n):
+        it = iter(iterable)
+        return iter(lambda: tuple(itertools.islice(it, n)), ())
+
+
+batched(range(3), 2)  # B911: 0

tests/eval_files/b911_shadowed_fallback.py

@@ -0,0 +1,17 @@
+"""
+Should emit:
+B911 - on line 17
+"""
+
+# A fallback for older Pythons still means batched can be itertools.batched.
+try:
+    from itertools import batched
+except ImportError:
+    import itertools
+
+    def batched(iterable, n):
+        it = iter(iterable)
+        return iter(lambda: tuple(itertools.islice(it, n)), ())
+
+
+batched(range(3), 2)  # B911: 0

tests/eval_files/b911_shadowed_import.py

@@ -0,0 +1,11 @@
+"""
+Should emit:
+B911 - on line 11
+"""
+
+import itertools
+
+from mylib import batched
+
+batched(range(3), 2)
+itertools.batched(range(3), 2)  # B911: 0