Skip to content

Commit fe09b02

Browse files
authored
Use match for statistics brain (#2835)
1 parent 2bca4f5 commit fe09b02

File tree

1 file changed

+25
-30
lines changed

1 file changed

+25
-30
lines changed

astroid/brain/brain_statistics.py

Lines changed: 25 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -13,49 +13,44 @@
1313
from collections.abc import Iterator
1414
from typing import TYPE_CHECKING
1515

16+
from astroid import nodes
1617
from astroid.context import InferenceContext
1718
from astroid.inference_tip import inference_tip
1819
from astroid.manager import AstroidManager
19-
from astroid.nodes.node_classes import Attribute, Call, ImportFrom, Name
2020
from astroid.util import Uninferable
2121

2222
if TYPE_CHECKING:
2323
from astroid.typing import InferenceResult
2424

2525

26-
def _looks_like_statistics_quantiles(node: Call) -> bool:
26+
def _looks_like_statistics_quantiles(node: nodes.Call) -> bool:
2727
"""Check if this is a call to statistics.quantiles."""
28-
# Case 1: statistics.quantiles(...)
29-
if isinstance(node.func, Attribute):
30-
if node.func.attrname != "quantiles":
31-
return False
32-
if isinstance(node.func.expr, Name):
33-
if node.func.expr.name == "statistics":
34-
return True
35-
36-
# Case 2: from statistics import quantiles; quantiles(...)
37-
if isinstance(node.func, Name) and node.func.name == "quantiles":
38-
# Check if quantiles was imported from statistics
39-
try:
40-
frame = node.frame()
41-
if "quantiles" in frame.locals:
42-
# Look for import from statistics
43-
for stmt in frame.body:
44-
if (
45-
isinstance(stmt, ImportFrom)
46-
and stmt.modname == "statistics"
47-
and any(name[0] == "quantiles" for name in stmt.names or [])
48-
):
49-
return True
50-
except (AttributeError, TypeError):
51-
# If we can't determine the import context, be conservative
52-
pass
53-
28+
match node.func:
29+
case nodes.Attribute(expr=nodes.Name(name="statistics"), attrname="quantiles"):
30+
# Case 1: statistics.quantiles(...)
31+
return True
32+
case nodes.Name(name="quantiles"):
33+
# Case 2: from statistics import quantiles; quantiles(...)
34+
# Check if quantiles was imported from statistics
35+
try:
36+
frame = node.frame()
37+
if "quantiles" in frame.locals:
38+
# Look for import from statistics
39+
for stmt in frame.body:
40+
if (
41+
isinstance(stmt, nodes.ImportFrom)
42+
and stmt.modname == "statistics"
43+
and any(name[0] == "quantiles" for name in stmt.names or [])
44+
):
45+
return True
46+
except (AttributeError, TypeError):
47+
# If we can't determine the import context, be conservative
48+
pass
5449
return False
5550

5651

5752
def infer_statistics_quantiles(
58-
node: Call, context: InferenceContext | None = None
53+
node: nodes.Call, context: InferenceContext | None = None
5954
) -> Iterator[InferenceResult]:
6055
"""Infer the result of statistics.quantiles() calls.
6156
@@ -72,7 +67,7 @@ def infer_statistics_quantiles(
7267
def register(manager: AstroidManager) -> None:
7368
"""Register statistics-specific inference improvements."""
7469
manager.register_transform(
75-
Call,
70+
nodes.Call,
7671
inference_tip(infer_statistics_quantiles),
7772
_looks_like_statistics_quantiles,
7873
)

0 commit comments

Comments
 (0)