Coverage for packages / griffelib / src / griffe / _internal / agents / nodes / assignments.py: 100.00%

1# This module contains utilities for extracting information from assignment nodes. 

2 

3from __future__ import annotations 

4 

5import ast 

6from typing import TYPE_CHECKING, Any 

7 

8if TYPE_CHECKING: 

9 from collections.abc import Callable 

10 

11 

12def _get_attribute_name(node: ast.Attribute) -> str: 

13 return f"{get_name(node.value)}.{node.attr}" 

14 

15 

16def _get_name_name(node: ast.Name) -> str: 

17 return node.id 

18 

19 

20_node_name_map: dict[type, Callable[[Any], str]] = { 

21 ast.Name: _get_name_name, 

22 ast.Attribute: _get_attribute_name, 

23} 

24 

25 

26def get_name(node: ast.AST) -> str: 

27 """Extract name from an assignment node. 

28 

29 Parameters: 

30 node: The node to extract names from. 

31 

32 Returns: 

33 A list of names. 

34 """ 

35 return _node_name_map[type(node)](node) 

36 

37 

38def _get_assign_names(node: ast.Assign) -> list[str]: 

39 names = (get_name(target) for target in node.targets) 

40 return [name for name in names if name] 

41 

42 

43def _get_annassign_names(node: ast.AnnAssign) -> list[str]: 

44 name = get_name(node.target) 

45 return [name] if name else [] 

46 

47 

48_node_names_map: dict[type, Callable[[Any], list[str]]] = { 

49 ast.Assign: _get_assign_names, 

50 ast.AnnAssign: _get_annassign_names, 

51} 

52 

53 

54def get_names(node: ast.AST) -> list[str]: 

55 """Extract names from an assignment node. 

56 

57 Parameters: 

58 node: The node to extract names from. 

59 

60 Returns: 

61 A list of names. 

62 """ 

63 return _node_names_map[type(node)](node) 

64 

65 

66def get_instance_names(node: ast.AST) -> list[str]: 

67 """Extract names from an assignment node, only for instance attributes. 

68 

69 Parameters: 

70 node: The node to extract names from. 

71 

72 Returns: 

73 A list of names. 

74 """ 

75 return [name.split(".", 1)[1] for name in get_names(node) if name.startswith("self.")]