Coverage for src/griffe_pydantic/extension.py: 49.18%

1"""Griffe extension for Pydantic.""" 

2 

3from __future__ import annotations 

4 

5import ast 

6from typing import TYPE_CHECKING, Any 

7 

8from griffe import ( 

9 Attribute, 

10 Class, 

11 Extension, 

12 Function, 

13 Module, 

14 get_logger, 

15) 

16 

17from griffe_pydantic import dynamic, static 

18 

19if TYPE_CHECKING: 

20 from griffe import ObjectNode 

21 

22 

23logger = get_logger(__name__) 

24 

25 

26class PydanticExtension(Extension): 

27 """Griffe extension for Pydantic.""" 

28 

29 def __init__(self, *, schema: bool = False) -> None: 

30 """Initialize the extension. 

31 

32 Parameters: 

33 schema: Whether to compute and store the JSON schema of models. 

34 """ 

35 super().__init__() 

36 self.schema = schema 

37 self.in_model: list[Class] = [] 

38 self.processed: set[str] = set() 

39 

40 def on_package_loaded(self, *, pkg: Module, **kwargs: Any) -> None: # noqa: ARG002 

41 """Detect models once the whole package is loaded.""" 

42 static.process_module(pkg, processed=self.processed, schema=self.schema) 

43 

44 def on_class_instance(self, *, node: ast.AST | ObjectNode, cls: Class, **kwargs: Any) -> None: # noqa: ARG002 

45 """Detect and prepare Pydantic models.""" 

46 # Prevent running during static analysis. 

47 if isinstance(node, ast.AST): 47 ↛ 50line 47 didn't jump to line 50 because the condition on line 47 was always true

48 return 

49 

50 try: 

51 import pydantic 

52 except ImportError: 

53 logger.warning("could not import pydantic - models will not be detected") 

54 return 

55 

56 if issubclass(node.obj, pydantic.BaseModel): 

57 self.in_model.append(cls) 

58 dynamic.process_class(node, cls) 

59 self.processed.add(cls.canonical_path) 

60 

61 def on_attribute_instance(self, *, node: ast.AST | ObjectNode, attr: Attribute, **kwargs: Any) -> None: # noqa: ARG002 

62 """Handle Pydantic fields.""" 

63 # Prevent running during static analysis. 

64 if isinstance(node, ast.AST): 64 ↛ 66line 64 didn't jump to line 66 because the condition on line 64 was always true

65 return 

66 if self.in_model: 

67 cls = self.in_model[-1] 

68 dynamic.process_attribute(node, attr, cls) 

69 self.processed.add(attr.canonical_path) 

70 

71 def on_function_instance(self, *, node: ast.AST | ObjectNode, func: Function, **kwargs: Any) -> None: # noqa: ARG002 

72 """Handle Pydantic field validators.""" 

73 # Prevent running during static analysis. 

74 if isinstance(node, ast.AST): 74 ↛ 76line 74 didn't jump to line 76 because the condition on line 74 was always true

75 return 

76 if self.in_model: 

77 cls = self.in_model[-1] 

78 dynamic.process_function(node, func, cls) 

79 self.processed.add(func.canonical_path) 

80 

81 def on_class_members(self, *, node: ast.AST | ObjectNode, cls: Class, **kwargs: Any) -> None: # noqa: ARG002 

82 """Finalize the Pydantic model data.""" 

83 # Prevent running during static analysis. 

84 if isinstance(node, ast.AST): 84 ↛ 87line 84 didn't jump to line 87 because the condition on line 84 was always true

85 return 

86 

87 if self.in_model and cls is self.in_model[-1]: 

88 # Pop last class from the heap. 

89 self.in_model.pop()