Coverage for packages / griffelib / src / griffe / _internal / agents / nodes / exports.py: 92.00%

1# This module contains utilities for extracting exports from `__all__` assignments. 

2 

3from __future__ import annotations 

4 

5import ast 

6from contextlib import suppress 

7from typing import TYPE_CHECKING, Any 

8 

9from griffe._internal.agents.nodes.values import get_value 

10from griffe._internal.enumerations import LogLevel 

11from griffe._internal.expressions import ExprName 

12from griffe._internal.logger import logger 

13 

14if TYPE_CHECKING: 

15 from collections.abc import Callable 

16 

17 from griffe._internal.models import Module 

18 

19 

20def _extract_attribute(node: ast.Attribute, parent: Module) -> list[str | ExprName]: 

21 return [ExprName(name=node.attr, parent=_extract(node.value, parent)[0])] 

22 

23 

24def _extract_binop(node: ast.BinOp, parent: Module) -> list[str | ExprName]: 

25 left = _extract(node.left, parent) 

26 right = _extract(node.right, parent) 

27 return left + right 

28 

29 

30def _extract_constant(node: ast.Constant, parent: Module) -> list[str | ExprName]: 

31 return [node.value] # ty:ignore[invalid-return-type] 

32 

33 

34def _extract_name(node: ast.Name, parent: Module) -> list[str | ExprName]: 

35 return [ExprName(node.id, parent)] 

36 

37 

38def _extract_sequence(node: ast.List | ast.Set | ast.Tuple, parent: Module) -> list[str | ExprName]: 

39 sequence = [] 

40 for elt in node.elts: 

41 sequence.extend(_extract(elt, parent)) 

42 return sequence 

43 

44 

45def _extract_starred(node: ast.Starred, parent: Module) -> list[str | ExprName]: 

46 return _extract(node.value, parent) 

47 

48 

49_node_map: dict[type, Callable[[Any, Module], list[str | ExprName]]] = { 

50 ast.Attribute: _extract_attribute, 

51 ast.BinOp: _extract_binop, 

52 ast.Constant: _extract_constant, 

53 ast.List: _extract_sequence, 

54 ast.Name: _extract_name, 

55 ast.Set: _extract_sequence, 

56 ast.Starred: _extract_starred, 

57 ast.Tuple: _extract_sequence, 

58} 

59 

60 

61def _extract(node: ast.AST, parent: Module) -> list[str | ExprName]: 

62 return _node_map[type(node)](node, parent) 

63 

64 

65def get__all__(node: ast.Assign | ast.AnnAssign | ast.AugAssign, parent: Module) -> list[str | ExprName]: 

66 """Get the values declared in `__all__`. 

67 

68 Parameters: 

69 node: The assignment node. 

70 parent: The parent module. 

71 

72 Returns: 

73 A set of names. 

74 """ 

75 if node.value is None: 75 ↛ 76line 75 didn't jump to line 76 because the condition on line 75 was never true

76 return [] 

77 return _extract(node.value, parent) 

78 

79 

80def safe_get__all__( 

81 node: ast.Assign | ast.AnnAssign | ast.AugAssign, 

82 parent: Module, 

83 log_level: LogLevel = LogLevel.debug, # TODO: Set to error when we handle more things? 

84) -> list[str | ExprName]: 

85 """Safely (no exception) extract values in `__all__`. 

86 

87 Parameters: 

88 node: The `__all__` assignment node. 

89 parent: The parent used to resolve the names. 

90 log_level: Log level to use to log a message. 

91 

92 Returns: 

93 A list of strings or resolvable names. 

94 """ 

95 try: 

96 return get__all__(node, parent) 

97 except Exception as error: # noqa: BLE001 

98 message = f"Failed to extract `__all__` value: {get_value(node.value)}" 

99 with suppress(Exception): 

100 message += f" at {parent.relative_filepath}:{node.lineno}" 

101 if isinstance(error, KeyError): 101 ↛ 104line 101 didn't jump to line 104 because the condition on line 101 was always true

102 message += f": unsupported node {error}" 

103 else: 

104 message += f": {error}" 

105 getattr(logger, log_level.value)(message) 

106 return []