Add brain module for statistics inference (#2832)

Author: emmanuel-ferdman

astroid/brain/brain_statistics.py

@@ -0,0 +1,78 @@
+# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
+# For details: https://github.com/pylint-dev/astroid/blob/main/LICENSE
+# Copyright (c) https://github.com/pylint-dev/astroid/blob/main/CONTRIBUTORS.txt
+
+"""Astroid hooks for understanding statistics library module.
+
+Provides inference improvements for statistics module functions that have
+complex runtime behavior difficult to analyze statically.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import TYPE_CHECKING
+
+from astroid.context import InferenceContext
+from astroid.inference_tip import inference_tip
+from astroid.manager import AstroidManager
+from astroid.nodes.node_classes import Attribute, Call, ImportFrom, Name
+from astroid.util import Uninferable
+
+if TYPE_CHECKING:
+    from astroid.typing import InferenceResult
+
+
+def _looks_like_statistics_quantiles(node: Call) -> bool:
+    """Check if this is a call to statistics.quantiles."""
+    # Case 1: statistics.quantiles(...)
+    if isinstance(node.func, Attribute):
+        if node.func.attrname != "quantiles":
+            return False
+        if isinstance(node.func.expr, Name):
+            if node.func.expr.name == "statistics":
+                return True
+
+    # Case 2: from statistics import quantiles; quantiles(...)
+    if isinstance(node.func, Name) and node.func.name == "quantiles":
+        # Check if quantiles was imported from statistics
+        try:
+            frame = node.frame()
+            if "quantiles" in frame.locals:
+                # Look for import from statistics
+                for stmt in frame.body:
+                    if (
+                        isinstance(stmt, ImportFrom)
+                        and stmt.modname == "statistics"
+                        and any(name[0] == "quantiles" for name in stmt.names or [])
+                    ):
+                        return True
+        except (AttributeError, TypeError):
+            # If we can't determine the import context, be conservative
+            pass
+
+    return False
+
+
+def infer_statistics_quantiles(
+    node: Call, context: InferenceContext | None = None
+) -> Iterator[InferenceResult]:
+    """Infer the result of statistics.quantiles() calls.
+
+    Returns Uninferable because quantiles() has complex runtime behavior
+    that cannot be statically analyzed, preventing false positives in
+    pylint's unbalanced-tuple-unpacking checker.
+
+    statistics.quantiles() returns a list with (n-1) elements, but static
+    analysis sees only the empty list initializations in the function body.
+    """
+    yield Uninferable
+
+
+def register(manager: AstroidManager) -> None:
+    """Register statistics-specific inference improvements."""
+    manager.register_transform(
+        Call,
+        inference_tip(infer_statistics_quantiles),
+        _looks_like_statistics_quantiles,
+    )

astroid/brain/helpers.py

@@ -75,6 +75,7 @@ def register_all_brains(manager: AstroidManager) -> None:
         brain_six,
         brain_sqlalchemy,
         brain_ssl,
+        brain_statistics,
         brain_subprocess,
         brain_threading,
         brain_type,
@@ -126,6 +127,7 @@ def register_all_brains(manager: AstroidManager) -> None:
     brain_six.register(manager)
     brain_sqlalchemy.register(manager)
     brain_ssl.register(manager)
+    brain_statistics.register(manager)
     brain_subprocess.register(manager)
     brain_threading.register(manager)
     brain_type.register(manager)

tests/brain/test_brain.py

@@ -321,6 +321,16 @@ class Derived(collections.abc.Hashable, collections.abc.Iterator[int]):
             ],
         )
 
+    def test_statistics_quantiles_from_import(self):
+        node = builder.extract_node(
+            """
+        from statistics import quantiles
+        quantiles([1, 2, 3, 4, 5, 6, 7, 8, 9], n=4)
+        """
+        )
+        inferred = next(node.infer())
+        self.assertIs(inferred, util.Uninferable)
+
 
 class TypingBrain(unittest.TestCase):
     def test_namedtuple_base(self) -> None:

tests/brain/test_statistics.py

@@ -0,0 +1,68 @@
+# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
+# For details: https://github.com/pylint-dev/astroid/blob/main/LICENSE
+# Copyright (c) https://github.com/pylint-dev/astroid/blob/main/CONTRIBUTORS.txt
+
+"""Tests for brain statistics module."""
+
+from __future__ import annotations
+
+import unittest
+
+from astroid import extract_node
+from astroid.util import Uninferable
+
+
+class StatisticsBrainTest(unittest.TestCase):
+    """Test the brain statistics module functionality."""
+
+    def test_statistics_quantiles_inference(self) -> None:
+        """Test that statistics.quantiles() returns Uninferable instead of empty list."""
+        node = extract_node(
+            """
+        import statistics
+        statistics.quantiles(list(range(100)), n=4)  #@
+        """
+        )
+        inferred = list(node.infer())
+        self.assertEqual(len(inferred), 1)
+        self.assertIs(inferred[0], Uninferable)
+
+    def test_statistics_quantiles_different_args(self) -> None:
+        """Test statistics.quantiles with different arguments."""
+        node = extract_node(
+            """
+        import statistics
+        statistics.quantiles([1, 2, 3, 4, 5], n=10, method='inclusive')  #@
+        """
+        )
+        inferred = list(node.infer())
+        self.assertEqual(len(inferred), 1)
+        self.assertIs(inferred[0], Uninferable)
+
+    def test_statistics_quantiles_assignment_unpacking(self) -> None:
+        """Test the specific case that was causing false positives."""
+        node = extract_node(
+            """
+        import statistics
+        q1, q2, q3 = statistics.quantiles(list(range(100)), n=4)  #@
+        """
+        )
+        call_node = node.value
+        inferred = list(call_node.infer())
+        self.assertEqual(len(inferred), 1)
+        self.assertIs(inferred[0], Uninferable)
+
+    def test_other_statistics_functions_not_affected(self) -> None:
+        """Test that other statistics functions are not affected by our brain module."""
+        node = extract_node(
+            """
+        import statistics
+        statistics.mean([1, 2, 3, 4, 5])  #@
+        """
+        )
+        inferred = list(node.infer())
+        self.assertNotEqual(len(inferred), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()