Coverage for src/griffe_pydantic/_internal/extension.py: 90.91%

1from __future__ import annotations 

2 

3import ast 

4from typing import TYPE_CHECKING, Any 

5 

6from griffe import ( 

7 Class, 

8 Extension, 

9 Module, 

10 get_logger, 

11) 

12 

13from griffe_pydantic._internal import dynamic, static 

14 

15if TYPE_CHECKING: 

16 from griffe import ObjectNode 

17 

18 

19_logger = get_logger("griffe_pydantic") 

20 

21 

22class PydanticExtension(Extension): 

23 """Griffe extension for Pydantic.""" 

24 

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

26 """Initialize the extension. 

27 

28 Parameters: 

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

30 """ 

31 super().__init__() 

32 self._schema = schema 

33 self._processed: set[str] = set() 

34 self._recorded: list[tuple[ObjectNode, Class]] = [] 

35 

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

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

38 for node, cls in self._recorded: 

39 self._processed.add(cls.canonical_path) 

40 dynamic._process_class(node.obj, cls, processed=self._processed, schema=self._schema) 

41 static._process_module(pkg, processed=self._processed, schema=self._schema) 

42 

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

44 """Detect and prepare Pydantic models.""" 

45 # Prevent running during static analysis. 

46 if isinstance(node, ast.AST): 

47 return 

48 

49 try: 

50 import pydantic # noqa: PLC0415 

51 except ImportError: 

52 _logger.warning("could not import pydantic - models will not be detected") 

53 return 

54 

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

56 self._recorded.append((node, cls))