Coverage for tests/test_extension.py: 100.00%

1"""Tests for the `extension` module.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6from typing import TYPE_CHECKING 

7 

8from griffe import Extensions, temporary_visited_package 

9 

10from griffe_pydantic.extension import PydanticExtension 

11 

12if TYPE_CHECKING: 

13 import pytest 

14 from mkdocstrings_handlers.python.handler import PythonHandler 

15 

16 

17code = """ 

18 from pydantic import field_validator, ConfigDict, BaseModel, Field 

19 

20 

21 class ExampleParentModel(BaseModel): 

22 '''An example parent model.''' 

23 parent_field: str = Field(..., description="Parent field.") 

24 

25 

26 class ExampleModel(ExampleParentModel): 

27 '''An example child model.''' 

28 

29 model_config = ConfigDict(frozen=False) 

30 

31 field_without_default: str 

32 '''Shows the *[Required]* marker in the signature.''' 

33 

34 field_plain_with_validator: int = 100 

35 '''Show standard field with type annotation.''' 

36 

37 field_with_validator_and_alias: str = Field("FooBar", alias="BarFoo", validation_alias="BarFoo") 

38 '''Shows corresponding validator with link/anchor.''' 

39 

40 field_with_constraints_and_description: int = Field( 

41 default=5, ge=0, le=100, description="Shows constraints within doc string." 

42 ) 

43 

44 @field_validator("field_with_validator_and_alias", "field_plain_with_validator", mode="before") 

45 @classmethod 

46 def check_max_length_ten(cls, v): 

47 '''Show corresponding field with link/anchor.''' 

48 if len(v) >= 10: 

49 raise ValueError("No more than 10 characters allowed") 

50 return v 

51 

52 def regular_method(self): 

53 pass 

54 

55 

56 class RegularClass(object): 

57 regular_attr = 1 

58""" 

59 

60 

61def test_extension() -> None: 

62 """Test the extension.""" 

63 with temporary_visited_package( 

64 "package", 

65 modules={"__init__.py": code}, 

66 extensions=Extensions(PydanticExtension(schema=True)), 

67 ) as package: 

68 assert package 

69 

70 assert "ExampleParentModel" in package.classes 

71 assert package.classes["ExampleParentModel"].labels == {"pydantic-model"} 

72 

73 assert "ExampleModel" in package.classes 

74 assert package.classes["ExampleModel"].labels == {"pydantic-model"} 

75 

76 config = package.classes["ExampleModel"].extra["griffe_pydantic"]["config"] 

77 assert config == {"frozen": "False"} 

78 

79 schema = package.classes["ExampleModel"].extra["griffe_pydantic"]["schema"] 

80 assert schema.startswith('{\n "description"') 

81 

82 

83def test_imported_models() -> None: 

84 """Test the extension with imported models.""" 

85 with temporary_visited_package( 

86 "package", 

87 modules={ 

88 "__init__.py": "from ._private import MyModel\n\n__all__ = ['MyModel']", 

89 "_private.py": "from pydantic import BaseModel\n\nclass MyModel(BaseModel):\n field1: str\n '''Some field.'''\n", 

90 }, 

91 extensions=Extensions(PydanticExtension(schema=False)), 

92 ) as package: 

93 assert package["MyModel"].labels == {"pydantic-model"} 

94 assert package["MyModel.field1"].labels == {"pydantic-field"} 

95 

96 

97def test_rendering_model_config_using_configdict(python_handler: PythonHandler) -> None: 

98 """Test the extension with model config using ConfigDict.""" 

99 code = """ 

100 from pydantic import BaseModel, ConfigDict, Field 

101 

102 class Model(BaseModel): 

103 usage: str | None = Field( 

104 None, 

105 description="Some description.", 

106 example="Some example.", 

107 ) 

108 model_config = ConfigDict( 

109 json_schema_extra={ 

110 "example": { 

111 "usage": "Some usage.", 

112 "limitations": "Some limitations.", 

113 "billing": "Some value.", 

114 "notice_period": "Some value.", 

115 } 

116 } 

117 ) 

118 """ 

119 with temporary_visited_package( 

120 "package", 

121 modules={"__init__.py": code}, 

122 extensions=Extensions(PydanticExtension(schema=False)), 

123 ) as package: 

124 python_handler.render(package["Model"], python_handler.get_options({})) # Assert no errors. 

125 

126 

127def test_not_crashing_on_dynamic_field_description(caplog: pytest.LogCaptureFixture) -> None: 

128 """Test the extension with dynamic field description.""" 

129 code = """ 

130 import pydantic 

131 

132 desc = "xyz" 

133 

134 class TestModel(pydantic.BaseModel): 

135 abc: str = pydantic.Field(description=desc) 

136 """ 

137 with ( 

138 caplog.at_level(logging.DEBUG), 

139 temporary_visited_package( 

140 "package", 

141 modules={"__init__.py": code}, 

142 extensions=Extensions(PydanticExtension(schema=False)), 

143 ), 

144 ): 

145 assert any( 

146 record.levelname == "DEBUG" and "field 'package.TestModel.abc' as literal" in record.message 

147 for record in caplog.records 

148 ) 

149 

150 

151def test_ignore_classvars() -> None: 

152 """Test the extension ignores class variables.""" 

153 code = """ 

154 from pydantic import BaseModel 

155 from typing import ClassVar 

156 

157 class Model(BaseModel): 

158 field: str 

159 class_var: ClassVar[int] = 1 

160 """ 

161 with temporary_visited_package( 

162 "package", 

163 modules={"__init__.py": code}, 

164 extensions=Extensions(PydanticExtension(schema=False)), 

165 ) as package: 

166 assert "pydantic-field" not in package["Model.class_var"].labels 

167 assert "class-attribute" in package["Model.class_var"].labels