Add subscribe method to AST nodes for stream initialization

With this subscription can be delegated to individual nodes, allowing
coalesce to subscribe to only the interesting stuff.

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/sdk/timeseries/formulas/_ast.py

@@ -61,10 +61,10 @@ def format(self, wrap: bool = False) -> str:
         """Return a string representation of the telemetry stream node."""
         return f"{self.source}"
 
-    async def fetch_next(self) -> None:
+    async def _fetch_next(self) -> None:
         """Fetch the next value for this component and store it internally."""
         if self._stream is None:
-            await self._fetch_stream()
+            await self.subscribe()
         assert self._stream is not None
 
         latest_sample = await anext(self._stream)
@@ -79,10 +79,12 @@ async def fetch_next(self) -> None:
         else:
             self._latest_sample = latest_sample
 
-    async def _fetch_stream(self) -> None:
+    @override
+    async def subscribe(self) -> None:
         """Subscribe to the telemetry stream for this component."""
         if self._stream is not None:
             return
+        _logger.debug("Subscribing to telemetry stream for %s", self.source)
         if self.metric_fetcher is None:
             raise RuntimeError("Metric fetcher is not set for TelemetryStream node.")
         self._stream = await self.metric_fetcher()
@@ -109,6 +111,11 @@ def format(self, wrap: bool = False) -> str:
         """Return a string representation of the function call node."""
         return self.function.format()
 
+    @override
+    async def subscribe(self) -> None:
+        """Subscribe to any data streams needed by the function."""
+        await self.function.subscribe()
+
 
 @dataclass(kw_only=True)
 class Constant(AstNode[QuantityT]):
@@ -126,6 +133,10 @@ def format(self, wrap: bool = False) -> str:
         """Return a string representation of the constant node."""
         return str(self.value.base_value)
 
+    @override
+    async def subscribe(self) -> None:
+        """Subscribe to any data streams needed by the function."""
+
 
 @dataclass(kw_only=True)
 class Add(AstNode[QuantityT]):
@@ -181,6 +192,14 @@ def format(self, wrap: bool = False) -> str:
             expr = f"({expr})"
         return expr
 
+    @override
+    async def subscribe(self) -> None:
+        """Subscribe to any data streams needed by the function."""
+        _ = await asyncio.gather(
+            self.left.subscribe(),
+            self.right.subscribe(),
+        )
+
 
 @dataclass(kw_only=True)
 class Sub(AstNode[QuantityT]):
@@ -237,6 +256,14 @@ def format(self, wrap: bool = False) -> str:
             expr = f"({expr})"
         return expr
 
+    @override
+    async def subscribe(self) -> None:
+        """Subscribe to any data streams needed by the function."""
+        _ = await asyncio.gather(
+            self.left.subscribe(),
+            self.right.subscribe(),
+        )
+
 
 @dataclass(kw_only=True)
 class Mul(AstNode[QuantityT]):
@@ -291,6 +318,14 @@ def format(self, wrap: bool = False) -> str:
         """Return a string representation of the multiplication node."""
         return f"{self.left.format(True)} * {self.right.format(True)}"
 
+    @override
+    async def subscribe(self) -> None:
+        """Subscribe to any data streams needed by the function."""
+        _ = await asyncio.gather(
+            self.left.subscribe(),
+            self.right.subscribe(),
+        )
+
 
 @dataclass(kw_only=True)
 class Div(AstNode[QuantityT]):
@@ -343,3 +378,11 @@ async def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
     def format(self, wrap: bool = False) -> str:
         """Return a string representation of the division node."""
         return f"{self.left.format(True)} / {self.right.format(True)}"
+
+    @override
+    async def subscribe(self) -> None:
+        """Subscribe to any data streams needed by the function."""
+        _ = await asyncio.gather(
+            self.left.subscribe(),
+            self.right.subscribe(),
+        )

src/frequenz/sdk/timeseries/formulas/_base_ast_node.py

@@ -31,3 +31,7 @@ def format(self, wrap: bool = False) -> str:
     def __str__(self) -> str:
         """Return the string representation of the node."""
         return self.format()
+
+    @abc.abstractmethod
+    async def subscribe(self) -> None:
+        """Subscribe to any data streams needed by this node."""

src/frequenz/sdk/timeseries/formulas/_functions.py

@@ -37,6 +37,12 @@ def format(self) -> str:
         params_str = ", ".join(str(param) for param in self.params)
         return f"{self.name}({params_str})"
 
+    async def subscribe(self) -> None:
+        """Subscribe to any data streams needed by the function."""
+        _ = await asyncio.gather(
+            *(param.subscribe() for param in self.params),
+        )
+
     @classmethod
     def from_string(
         cls, name: str, params: list[AstNode[QuantityT]]