Coverage for src/_griffe/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 Any, Callable 

7 

8 

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

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

11 

12 

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

14 return node.id 

15 

16 

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

18 ast.Name: _get_name_name, 

19 ast.Attribute: _get_attribute_name, 

20} 

21 

22 

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

24 """Extract name from an assignment node. 

25 

26 Parameters: 

27 node: The node to extract names from. 

28 

29 Returns: 

30 A list of names. 

31 """ 

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

33 

34 

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

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

37 return [name for name in names if name] 

38 

39 

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

41 name = get_name(node.target) 

42 return [name] if name else [] 

43 

44 

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

46 ast.Assign: _get_assign_names, 

47 ast.AnnAssign: _get_annassign_names, 

48} 

49 

50 

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

52 """Extract names from an assignment node. 

53 

54 Parameters: 

55 node: The node to extract names from. 

56 

57 Returns: 

58 A list of names. 

59 """ 

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

61 

62 

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

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

65 

66 Parameters: 

67 node: The node to extract names from. 

68 

69 Returns: 

70 A list of names. 

71 """ 

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