Coverage for .nox/test-3-9/lib/python3.9/site-packages/nskit/common/configuration/__init__.py: 97%

1"""Base configuration class. 

2 

3Includes: 

4 - properties in model dumps 

5 - file based config loading (json/toml/yaml) 

6 - model dumping to toml & yaml 

7""" 

8from __future__ import annotations 

9 

10from pathlib import Path 

11from typing import Any, Optional 

12 

13from pydantic.config import ExtraValues 

14from pydantic_settings import ( 

15 BaseSettings as _BaseSettings, 

16 PydanticBaseSettingsSource as _PydanticBaseSettingsSource, 

17 SettingsConfigDict as _SettingsConfigDict, 

18) 

19from pydantic_settings.sources import PathType 

20 

21from nskit.common.configuration.mixins import PropertyDumpMixin 

22from nskit.common.configuration.sources import ( 

23 DotEnvSettingsSource, 

24 JsonConfigSettingsSource, 

25 TomlConfigSettingsSource, 

26 YamlConfigSettingsSource, 

27) 

28from nskit.common.io import json, toml, yaml 

29 

30 

31class SettingsConfigDict(_SettingsConfigDict): 

32 """Customised Settings Config Dict.""" 

33 dotenv_extra: Optional[ExtraValues] = 'ignore' 

34 config_file: Optional[PathType] = None 

35 config_file_encoding: Optional[str] = None 

36 

37 

38class BaseConfiguration(PropertyDumpMixin, _BaseSettings): 

39 """A Pydantic BaseSettings type object with Properties included in model dump, and yaml and toml integrations.""" 

40 

41 model_config = SettingsConfigDict(env_file_encoding='utf-8') 

42 

43 def settings_customise_sources( 

44 cls, 

45 settings_cls: type[_BaseSettings], 

46 init_settings: _PydanticBaseSettingsSource, 

47 env_settings: _PydanticBaseSettingsSource, 

48 dotenv_settings: _PydanticBaseSettingsSource, 

49 file_secret_settings: _PydanticBaseSettingsSource, 

50 ) -> tuple[_PydanticBaseSettingsSource, ...]: 

51 """Create settings loading, including the FileConfigSettingsSource.""" 

52 config_files = cls.model_config.get('config_file') 

53 config_file_encoding = cls.model_config.get('config_file_encoding') 

54 file_types = {'json' : ['.json', '.jsn'], 

55 'yaml': ['.yaml', '.yml'], 

56 'toml': ['.toml', '.tml']} 

57 if config_files: 

58 if isinstance(config_files, (Path, str)): 

59 config_files = [config_files] 

60 else: 

61 config_files = [] 

62 

63 split_config_files = {} 

64 for file_type, suffixes in file_types.items(): 

65 original = cls.model_config.get(f'{file_type}_file') 

66 if original and isinstance(original, (Path, str)): 

67 split_config_files[file_type] = [original] 

68 elif original: 

69 split_config_files[file_type] = original 

70 else: 

71 split_config_files[file_type] = [] 

72 for config_file in config_files: 

73 if Path(config_file).suffix.lower() in suffixes: 

74 split_config_files[file_type].append(config_file) 

75 return ( 

76 init_settings, 

77 env_settings, 

78 JsonConfigSettingsSource(settings_cls, 

79 split_config_files['json'], 

80 cls.model_config.get('json_file_encoding') or config_file_encoding), 

81 YamlConfigSettingsSource(settings_cls, 

82 split_config_files['yaml'], 

83 cls.model_config.get('yaml_file_encoding') or config_file_encoding), 

84 TomlConfigSettingsSource(settings_cls, 

85 split_config_files['toml']), 

86 DotEnvSettingsSource(settings_cls, 

87 dotenv_settings.env_file, 

88 dotenv_settings.env_file_encoding, 

89 dotenv_settings.case_sensitive, 

90 dotenv_settings.env_prefix, 

91 dotenv_settings.env_nested_delimiter, 

92 dotenv_settings.env_ignore_empty, 

93 dotenv_settings.env_parse_none_str, 

94 cls.model_config.get('dotenv_extra', 'ignore')), 

95 file_secret_settings 

96 ) 

97 

98 def model_dump_toml( 

99 self, 

100 *, 

101 indent: int | None = None, 

102 include: Any = None, 

103 exclude: Any = None, 

104 by_alias: bool = False, 

105 exclude_unset: bool = False, 

106 exclude_defaults: bool = False, 

107 exclude_none: bool = False, 

108 round_trip: bool = False, 

109 warnings: bool = True): 

110 """Dump model to TOML.""" 

111 # We go via JSON to include indent etc. 

112 return toml.dumps(json.loads(self.model_dump_json( 

113 indent=indent, 

114 include=include, 

115 exclude=exclude, 

116 by_alias=by_alias, 

117 exclude_unset=exclude_unset, 

118 exclude_defaults=exclude_defaults, 

119 exclude_none=exclude_none, 

120 round_trip=round_trip, 

121 warnings=warnings 

122 ))) 

123 

124 def model_dump_yaml( 

125 self, 

126 *, 

127 indent: int | None = None, 

128 include: Any = None, 

129 exclude: Any = None, 

130 by_alias: bool = False, 

131 exclude_unset: bool = False, 

132 exclude_defaults: bool = False, 

133 exclude_none: bool = False, 

134 round_trip: bool = False, 

135 warnings: bool = True): 

136 """Dump model to YAML.""" 

137 # We go via JSON to include indent etc. 

138 return yaml.dumps(json.loads(self.model_dump_json( 

139 indent=indent, 

140 include=include, 

141 exclude=exclude, 

142 by_alias=by_alias, 

143 exclude_unset=exclude_unset, 

144 exclude_defaults=exclude_defaults, 

145 exclude_none=exclude_none, 

146 round_trip=round_trip, 

147 warnings=warnings 

148 )))