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

1"""Add settings sources.""" 

2from __future__ import annotations as _annotations 

3 

4from pathlib import Path 

5from typing import Any 

6 

7from pydantic.config import ExtraValues 

8from pydantic_settings import BaseSettings 

9from pydantic_settings.sources import ( 

10 DotEnvSettingsSource as _DotEnvSettingsSource, 

11 DotenvType, 

12 ENV_FILE_SENTINEL, 

13 JsonConfigSettingsSource as _JsonConfigSettingsSource, 

14 TomlConfigSettingsSource as _TomlConfigSettingsSource, 

15 YamlConfigSettingsSource as _YamlConfigSettingsSource, 

16) 

17 

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

19 

20 

21class JsonConfigSettingsSource(_JsonConfigSettingsSource): 

22 """Use the nskit.common.io.json loading to load settings from a json file.""" 

23 def _read_file(self, file_path: Path) -> dict[str, Any]: 

24 encoding = self.json_file_encoding or 'utf-8' 

25 file_contents = file_path.read_text(encoding) 

26 return json.loads(file_contents) 

27 

28 def __call__(self): 

29 """Make the file reading at the source instantiation.""" 

30 self.init_kwargs = self._read_files(self.json_file_path) 

31 return super().__call__() 

32 

33 

34class TomlConfigSettingsSource(_TomlConfigSettingsSource): 

35 """Use the nskit.common.io.toml loading to load settings from a toml file.""" 

36 def _read_file(self, file_path: Path) -> dict[str, Any]: 

37 file_contents = file_path.read_text() 

38 return toml.loads(file_contents) 

39 

40 def __call__(self): 

41 """Make the file reading at the source instantiation.""" 

42 self.init_kwargs = self._read_files(self.toml_file_path) 

43 return super().__call__() 

44 

45 

46class YamlConfigSettingsSource(_YamlConfigSettingsSource): 

47 """Use the nskit.common.io.yaml loading to load settings from a yaml file.""" 

48 def _read_file(self, file_path: Path) -> dict[str, Any]: 

49 encoding = self.yaml_file_encoding or 'utf-8' 

50 file_contents = file_path.read_text(encoding) 

51 return yaml.loads(file_contents) 

52 

53 def __call__(self): 

54 """Make the file reading at the source instantiation.""" 

55 self.init_kwargs = self._read_files(self.yaml_file_path) 

56 return super().__call__() 

57 

58 

59class DotEnvSettingsSource(_DotEnvSettingsSource): 

60 """Fixes change of behaviour in pydantic-settings 2.2.0 with extra allowed handling. 

61 

62 Adds dotenv_extra variable that is set to replicate previous behaviour (ignore). 

63 """ 

64 

65 def __init__( 

66 self, 

67 settings_cls: type[BaseSettings], 

68 env_file: DotenvType | None = ENV_FILE_SENTINEL, 

69 env_file_encoding: str | None = None, 

70 case_sensitive: bool | None = None, 

71 env_prefix: str | None = None, 

72 env_nested_delimiter: str | None = None, 

73 env_ignore_empty: bool | None = None, 

74 env_parse_none_str: str | None = None, 

75 dotenv_extra: ExtraValues | None = 'ignore' 

76 ) -> None: 

77 """Wrapper for init function to add dotenv_extra handling.""" 

78 self.dotenv_extra = dotenv_extra 

79 super().__init__( 

80 settings_cls, 

81 env_file, 

82 env_file_encoding, 

83 case_sensitive, 

84 env_prefix, 

85 env_nested_delimiter, 

86 env_ignore_empty, 

87 env_parse_none_str 

88 ) 

89 

90 def __call__(self) -> dict[str, Any]: 

91 """Wraps call logic introduced in 2.2.0, but is backwards compatible to 2.1.0 and earlier versions.""" 

92 data: dict[str, Any] = super().__call__() 

93 to_pop = [] 

94 for key in data.keys(): 

95 matched = False 

96 for field_name, field in self.settings_cls.model_fields.items(): 

97 for field_alias, field_env_name, _ in self._extract_field_info(field, field_name): 

98 if key == field_env_name or key == field_alias: 

99 matched = True 

100 break 

101 if matched: 

102 break 

103 if not matched and self.dotenv_extra == 'ignore': 

104 to_pop.append(key) 

105 for key in to_pop: 

106 data.pop(key) 

107 return data