Coverage for structured_tutorials / models / base.py: 100%

1# Copyright (c) 2025 Mathias Ertl 

2# Licensed under the MIT License. See LICENSE file for details. 

3 

4"""Base model classes.""" 

5 

6from pathlib import Path 

7from typing import Annotated, Any, Generic, TypeVar, overload 

8 

9from pydantic import ( 

10 BaseModel, 

11 ConfigDict, 

12 Field, 

13 NonNegativeFloat, 

14 NonNegativeInt, 

15 field_validator, 

16 model_validator, 

17) 

18from pydantic.fields import FieldInfo 

19 

20from structured_tutorials.typing import Self 

21 

22# Type for commands to execute 

23CommandType = str | tuple[str, ...] 

24V = TypeVar("V") 

25 

26TEMPLATE_DESCRIPTION = "This value is rendered as a template with the current context." 

27 

28 

29def template_field_title_generator(field_name: str, field_info: FieldInfo) -> str: 

30 """Field title generator for template fields.""" 

31 return f"{field_name.title()} (template)" 

32 

33 

34class CommandBaseModel(BaseModel): 

35 """Base model for commands.""" 

36 

37 model_config = ConfigDict(extra="forbid") 

38 

39 status_code: Annotated[int, Field(ge=0, le=255)] = 0 

40 clear_environment: bool = Field(default=False, description="Clear the environment.") 

41 update_environment: dict[str, str] = Field( 

42 default_factory=dict, description="Update the environment for all subsequent commands." 

43 ) 

44 environment: dict[str, str] = Field( 

45 default_factory=dict, description="Additional environment variables for the process." 

46 ) 

47 show_output: bool = Field( 

48 default=True, description="Set to `False` to always hide the output of this command." 

49 ) 

50 

51 

52class TestSpecificationMixin: 

53 """Mixin for specifying tests.""" 

54 

55 delay: Annotated[float, Field(ge=0)] = 0 

56 retry: NonNegativeInt = 0 

57 backoff_factor: NonNegativeFloat = 0 # {backoff factor} * (2 ** ({number of previous retries})) 

58 

59 

60class ConfigurationMixin: 

61 """Mixin for configuration models.""" 

62 

63 update_context: dict[str, Any] = Field(default_factory=dict) 

64 

65 

66class DocumentationConfigurationMixin: 

67 """Mixin for documentation configuration models.""" 

68 

69 text_before: str = Field(default="", description="Text before documenting this part.") 

70 text_after: str = Field(default="", description="Text after documenting this part.") 

71 

72 

73class DictRootModelMixin(Generic[V]): 

74 """Mixin for dict-based root models.""" 

75 

76 root: dict[str, V] 

77 

78 @overload 

79 def get(self, key: str, default: V) -> V: ... 

80 

81 @overload 

82 def get(self, key: str) -> V | None: ... 

83 

84 def get(self, key: str, default: V | None = None) -> V | None: 

85 return self.root.get(key, default) 

86 

87 

88class FileMixin: 

89 """Mixin for specifying a file (used in file part and for stdin of commands).""" 

90 

91 contents: str | None = Field( 

92 default=None, 

93 field_title_generator=template_field_title_generator, 

94 description=f"Contents of the file. {TEMPLATE_DESCRIPTION}", 

95 ) 

96 source: Path | None = Field( 

97 default=None, 

98 field_title_generator=template_field_title_generator, 

99 description="The source path of the file to create. Unless `template` is `False`, the file is loaded " 

100 "into memory and rendered as template.", 

101 ) 

102 template: bool = Field( 

103 default=True, description="Whether the file contents should be rendered in a template." 

104 ) 

105 

106 @field_validator("source", mode="after") 

107 @classmethod 

108 def validate_source(cls, value: Path) -> Path: 

109 if value.is_absolute(): 

110 raise ValueError(f"{value}: Must be a relative path (relative to the current cwd).") 

111 return value 

112 

113 @model_validator(mode="after") 

114 def validate_contents_or_source(self) -> Self: 

115 if self.contents is None and self.source is None: 

116 raise ValueError("Either contents or source is required.") 

117 if self.contents is not None and self.source is not None: 

118 raise ValueError("Only one of contents or source is allowed.") 

119 return self