Coverage for core / src / sensorkit / api / bootstrap.py: 39%

1# SPDX-License-Identifier: Apache-2.0 

2import asyncio 

3import os 

4import pathlib 

5import warnings 

6from collections.abc import Iterable 

7from typing import Literal, cast, overload 

8 

9from sensorkit.backend.base import BackendImpl 

10from sensorkit.common.importutil import import_module_or_file, obj_from_spec 

11from sensorkit.config.parser import SensorKitBaseConfig, SensorKitConfig, parse_config 

12 

13DEFAULT_BACKEND = "nats" 

14DEFAULT_CONFIG_FILE = "sensorkit.yaml" 

15DEFAULT_BASE_IMPORTS = ( 

16 "sensorkit.std", 

17 "sensorkit.webapi", 

18 "sensorkit.data.filesys", 

19 "sensorkit.data.fits", 

20 "sensorkit.data.local", 

21) 

22BACKEND_MODULES = { 

23 "nats": "sensorkit.backend.nats", 

24 "fake": "sensorkit.backend.fake", 

25} 

26 

27_read_lock = asyncio.Lock() 

28_base: SensorKitBaseConfig | None = None 

29_resolved: SensorKitConfig | None = None 

30 

31 

32@overload 

33async def _read_config(*, required: Literal[False] = False) -> SensorKitBaseConfig | None: ... 

34 

35 

36@overload 

37async def _read_config(*, required: Literal[True]) -> SensorKitBaseConfig: ... 

38 

39 

40async def _read_config(*, required: bool = False): 

41 """Read and memoize the base config. 

42 

43 Callers must hold `_read_lock`. 

44 """ 

45 global _base 

46 

47 def _read_config_sync(location: str | None): 

48 import yaml 

49 

50 path = pathlib.Path(location or DEFAULT_CONFIG_FILE) 

51 

52 try: 

53 base = parse_config(yaml.safe_load(path.read_text())) 

54 except FileNotFoundError: 

55 if required or location is not None: 

56 raise 

57 

58 base = None 

59 

60 return base 

61 

62 if _base is None: 

63 _base = await asyncio.to_thread( 

64 _read_config_sync, 

65 os.environ.get("SENSORKIT_CONFIG"), 

66 ) 

67 

68 return _base 

69 

70 

71def _imports_from_env(var: str, default: Iterable[str] = ()) -> list[str]: 

72 return [mod.strip() for mod in os.environ.get(var, "").split(",") if mod] or list(default) 

73 

74 

75def _import_modules_sync(*, fail_policy: Literal["error", "warn", "ignore"]): 

76 imports = _imports_from_env("SENSORKIT_BASE_IMPORTS", DEFAULT_BASE_IMPORTS) 

77 imports.extend(_imports_from_env("SENSORKIT_IMPORTS")) 

78 

79 if _base is not None: 

80 imports.extend(_base.configured_imports()) 

81 

82 for module in imports: 

83 try: 

84 import_module_or_file(module) 

85 except Exception: 

86 match fail_policy: 

87 case "error": 

88 raise 

89 case "warn": 

90 warnings.warn(f"Failed to import: {module}", stacklevel=1) 

91 case "ignore": 

92 pass 

93 

94 

95def set_config_location(location: str | None) -> None: 

96 """Set the location of the SensorKit configuration file. 

97 

98 This must be called before `load_config`, `import_modules`, and `connect` 

99 to have effect. 

100 """ 

101 if location is not None: 

102 current = os.environ.get("SENSORKIT_CONFIG", DEFAULT_CONFIG_FILE) 

103 

104 if _base is not None and location != current: 

105 raise RuntimeError(f"config already loaded from {current}, cannot set to {location}") 

106 

107 os.environ["SENSORKIT_CONFIG"] = location 

108 

109 

110async def load_config( 

111 *, 

112 fail_policy: Literal["error", "warn", "ignore"] = "warn", 

113) -> SensorKitConfig: 

114 """Loads and fully resolves configuration, importing modules as needed. 

115 

116 The fully resolved configuration is memoized; subsequent calls return the 

117 cached result without re-resolving. 

118 """ 

119 global _resolved 

120 

121 async with _read_lock: 

122 if _resolved is None: 

123 base = await _read_config(required=True) 

124 await asyncio.to_thread(_import_modules_sync, fail_policy=fail_policy) 

125 _resolved = base.resolve_dynamic_sections() 

126 

127 return cast(SensorKitConfig, _resolved) 

128 

129 

130async def import_modules(*, fail_policy: Literal["error", "warn", "ignore"] = "error"): 

131 """Imports modules based on environment and configuration.""" 

132 async with _read_lock: 

133 await _read_config() 

134 

135 await asyncio.to_thread(_import_modules_sync, fail_policy=fail_policy) 

136 

137 

138async def connect(*, backend: str | None = None): 

139 """Creates a SensorKit client and connects to the backend. 

140 

141 The backend is either provided as a string, read from the environment, or 

142 (if neither is present) read from the SensorKit configuration file. 

143 

144 User modules are not imported automatically. Callers that need to access 

145 types provided by modules must opt in to these imports by calling 

146 `import_modules()` or `load_config()` (to retrieve the fully resolved 

147 configuration structure). 

148 """ 

149 from sensorkit.core.client import SensorKit 

150 

151 backend = backend or os.environ.get("SENSORKIT_BACKEND") 

152 

153 if not backend: 

154 async with _read_lock: 

155 config = await _read_config() 

156 

157 backend = config.sensorkit.backend if config else None 

158 

159 backend_module = BACKEND_MODULES.get(backend) or BACKEND_MODULES[DEFAULT_BACKEND] 

160 backend_cls = await asyncio.to_thread( 

161 obj_from_spec, 

162 spec=backend_module, 

163 base=BackendImpl, 

164 subclass=True, 

165 ) 

166 backend_impl = await backend_cls.create() 

167 return SensorKit(backend=backend_impl)