Coverage for core / src / sensorkit / data / context.py: 96%
113 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-02 00:03 +0000
1# SPDX-License-Identifier: Apache-2.0
2from __future__ import annotations
4import asyncio
5import builtins
6from datetime import UTC, datetime
7from types import MappingProxyType
8from typing import TYPE_CHECKING, Any, Unpack, overload, override
10from pydantic import BaseModel
12from sensorkit.common.keyword import CompositeKeyword, KeywordDict
14if TYPE_CHECKING:
15 from sensorkit.core.entity import EntityClient
17_SAFE_BUILTINS = {
18 "__import__": builtins.__import__,
19 "format": format,
20 "str": str,
21 "int": int,
22 "float": float,
23 "bool": bool,
24 "round": round,
25 "len": len,
26 "min": min,
27 "max": max,
28 "abs": abs,
29}
31_EVAL_GLOBALS = {"__builtins__": _SAFE_BUILTINS, "datetime": datetime, "UTC": UTC}
32_MISSING = object()
33_FSTRING_OPENERS = ('f"', "f'", 'F"', "F'")
36def _raw_fstring_source(value: str) -> str:
37 """Build source that evaluates `value` as a raw f-string template.
39 The value is embedded in a triple-double-quoted raw f-string, so backslashes stay
40 literal while `{...}` still interpolates. Trailing backslashes or quotes cannot safely
41 abut the closing delimiter (a raw literal cannot end in an odd backslash run, and a
42 trailing quote would merge with it), so they are peeled off and re-appended verbatim
43 via concatenation. A triple-quote embedded mid-value is unsupported and surfaces as a
44 `SyntaxError` when the source is evaluated.
46 Args:
47 value: Template text containing one or more `{...}` fields.
49 Returns:
50 Python source that evaluates to the rendered string.
51 """
52 end = len(value)
54 while end and value[end - 1] in '\\"':
55 end -= 1
57 source = f'rf"""{value[:end]}"""'
59 if end < len(value):
60 source += f" + {value[end:]!r}"
62 return source
65class Context(KeywordDict):
66 """A `KeywordDict` read view that expands composite keywords.
68 A `KeywordDict` stores only what is explicitly set, so its serialized form stays minimal.
69 A `Context` is the resolvable view over that data: as keywords enter it, any
70 `CompositeKeyword` also makes its composed keywords available under their own keys, so a
71 consumer can look up (and `eval` against) a composed keyword without knowing which
72 composite carried it. Expansion happens eagerly on every write path — construction,
73 `set`, and `update` — because the data pipeline mutates a context in flight, so a
74 one-shot expand at construction would not hold.
76 A `Context` source is trusted to already be expanded; any other source (a bare
77 `KeywordDict`, an iterable of items from deserialization, keyword arguments) is expanded on
78 the way in.
79 """
81 def __init__(
82 self,
83 arg: Any = None,
84 *objs: Unpack[tuple[object, ...]],
85 ):
86 super().__init__(arg, *objs)
88 if not isinstance(arg, Context):
89 self._expand_all()
91 @override
92 def update(self, other):
93 super().update(other)
95 # Expand only what `other` just contributed.
96 if not isinstance(other, Context):
97 self._expand_values(other.values())
99 @override
100 def set_keyword(self, obj):
101 if isinstance(obj, CompositeKeyword):
102 self._expand(obj, {id(obj)})
104 super().set_keyword(obj)
106 def _expand_all(self):
107 self._expand_values(self.values())
109 def _expand_values(self, values):
110 for value in list(values):
111 if isinstance(value, CompositeKeyword):
112 self._expand(value, {id(value)})
114 def _expand(self, obj: CompositeKeyword, visited: set[int]):
115 # DFS over composed keywords, storing each under its own key. `super()._set_keyword`
116 # keeps the walk explicit and avoids re-entering this class's expanding `_set_keyword`.
117 for composed in obj.composed_keywords():
118 if id(composed) in visited:
119 continue
121 visited.add(id(composed))
123 if isinstance(composed, CompositeKeyword):
124 self._expand(composed, visited)
126 super().set_keyword(composed)
128 def eval(self, expr: str, *, default: object = _MISSING):
129 """Evaluate a Python expression with this Context as the namespace.
131 Names in the expression resolve against this Context's keys.
133 Args:
134 expr: The Python expression to evaluate.
135 default: Value returned if `expr` references a name absent from the Context.
136 If omitted, a missing name raises `NameError`. Only missing names are
137 caught; any other error (such as a bad format spec on a present value)
138 always propagates.
140 Returns:
141 The expression result, keeping its native type.
143 Raises:
144 NameError: A referenced name is absent and no `default` was given.
145 """
146 try:
147 return eval(expr, _EVAL_GLOBALS, self)
148 except NameError:
149 if default is _MISSING:
150 raise
152 return default
154 @overload
155 def resolve[T](self, value: str | None, *, as_type: type[T], default: T = ...) -> T: ...
157 @overload
158 def resolve(self, value: str | None, *, default: object = ...) -> object: ...
160 def resolve(self, value: str | None, *, as_type = _MISSING, default = _MISSING):
161 r"""Resolve a config string against this Context.
163 The form of `value` selects how it is interpreted:
165 - `=<expr>`: evaluate the remainder as a Python expression (see `eval`), keeping
166 its native type, e.g. `=FileInfo.path.name.upper()` or `=frame_num + 1`.
167 - `f"..."` or `F"..."`: evaluate as that f-string, yielding a string. Use this form
168 when escape sequences should be processed, e.g. `f"line\n{frame_num}"`.
169 - text containing `{...}`: evaluate as a raw f-string template, so fields
170 interpolate with full expression power while backslashes stay literal, e.g.
171 `C:\Temp\{frame_num}.fits`. Write `{{` or `}}` for a literal brace.
172 - anything else: literal text returned verbatim, e.g. `C:\Temp`.
174 Args:
175 value: The config string to resolve.
176 as_type: If given, raise TypeError if the resolved value is not of this type.
177 default: The resolved value if a referenced name is absent (see `eval`) or if
178 the input value is None.
180 Returns:
181 The resolved value: native type for `=<expr>`, a string for the f-string and
182 template forms, or the original text for a literal.
184 Raises:
185 TypeError: if `as_type` is given and the resolved value is not of that type.
186 """
187 if value is None:
188 if default is _MISSING:
189 raise TypeError("cannot resolve None without default")
191 value = default
192 elif value.startswith("="):
193 value = self.eval(value[1:], default=default)
194 elif value.startswith(_FSTRING_OPENERS):
195 value = self.eval(value, default=default)
196 elif "{" in value:
197 value = self.eval(_raw_fstring_source(value), default=default)
199 if as_type is not _MISSING and not isinstance(value, as_type):
200 raise TypeError(f"expected {as_type.__name__}, got {type(value).__name__}")
202 return value
205class ContextSubscription:
206 """Subscribe to device keyword updates and produce Context snapshots.
208 Monitors one or more keyword types on EntityClients, caching the latest
209 value of each. A `snapshot` merges the cached keyword models with an
210 optional base context and additional key-value pairs.
212 The cached keyword models are stored as-is in the context, preserving
213 the strict typing afforded by Keywords.
215 Example usage::
217 sub = ContextSubscription(mount_client)
218 sub.add(AltAzPointing)
219 sub.add(RADecPointing)
221 await sub.start()
222 ctx = sub.snapshot(task.execution.get_context())
223 await sub.stop()
224 """
226 def __init__(self, client: EntityClient):
227 self._client = client
228 self._add_queue: asyncio.Queue[type[BaseModel]] = asyncio.Queue()
229 self._add_task: asyncio.Task | None = None
230 self._consumers: list[asyncio.Task] = []
231 self._ready = asyncio.Event()
232 self._cache = KeywordDict()
233 self.cache = MappingProxyType(self._cache)
235 def add(self, keyword_type: type[BaseModel]):
236 """Register a keyword subscription.
238 Args:
239 keyword_type: The keyword/model type to subscribe to.
240 """
241 self._ready.clear()
242 self._add_queue.put_nowait(keyword_type)
244 async def _subscription_adder(self):
245 while True:
246 keyword = await self._add_queue.get()
248 if self._add_queue.empty():
249 self._ready.set()
251 # TODO: Need a monitor variant that allows a Queue parameter, then we only need one.
252 self._consumers.append(asyncio.create_task(self._consumer(keyword)))
254 async def _consumer(self, keyword: type[BaseModel]):
255 stream = await self._client.monitor(keyword)
257 async for _, data in stream:
258 self._cache.set(data)
260 async def start(self):
261 """Start all subscriptions."""
262 if self._add_task:
263 raise RuntimeError("ContextSubscription is already running")
265 self._add_task = asyncio.create_task(self._subscription_adder())
267 if self._add_queue.empty():
268 self._ready.set()
269 else:
270 await self._ready.wait()
272 async def stop(self):
273 """Cancel all background monitor tasks."""
274 if not self._add_task:
275 return
277 self._add_task.cancel()
279 for task in self._consumers:
280 task.cancel()
282 await asyncio.gather(self._add_task, *self._consumers, return_exceptions=True)
284 self._add_task = None
285 self._consumers.clear()
287 def snapshot(self, into: KeywordDict | None = None) -> Context:
288 """Copy the cached Context into a target Context."""
289 if into is None:
290 into = Context()
292 into.update(self._cache)
293 return into