Add an unsubscribe method to AST nodes

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

Author: shsms

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

@@ -89,6 +89,15 @@ async def subscribe(self) -> None:
             raise RuntimeError("Metric fetcher is not set for TelemetryStream node.")
         self._stream = await self.metric_fetcher()
 
+    @override
+    async def unsubscribe(self) -> None:
+        """Unsubscribe from the telemetry stream for this component."""
+        if self._stream is None:
+            return
+        _logger.debug("Unsubscribing from telemetry stream for %s", self.source)
+        self._stream.close()
+        self._stream = None
+
     def _is_quantity_sample(
         self, sample: Sample[QuantityT] | Sample[Quantity]
     ) -> TypeIs[Sample[Quantity]]:
@@ -116,6 +125,11 @@ async def subscribe(self) -> None:
         """Subscribe to any data streams needed by the function."""
         await self.function.subscribe()
 
+    @override
+    async def unsubscribe(self) -> None:
+        """Unsubscribe from any data streams needed by the function."""
+        await self.function.unsubscribe()
+
 
 @dataclass(kw_only=True)
 class Constant(AstNode[QuantityT]):
@@ -135,7 +149,11 @@ def format(self, wrap: bool = False) -> str:
 
     @override
     async def subscribe(self) -> None:
-        """Subscribe to any data streams needed by the function."""
+        """No-op for constant node."""
+
+    @override
+    async def unsubscribe(self) -> None:
+        """No-op for constant node."""
 
 
 @dataclass(kw_only=True)
@@ -196,12 +214,20 @@ def format(self, wrap: bool = False) -> str:
 
     @override
     async def subscribe(self) -> None:
-        """Subscribe to any data streams needed by the function."""
+        """Subscribe to any data streams needed by this node."""
         _ = await asyncio.gather(
             self.left.subscribe(),
             self.right.subscribe(),
         )
 
+    @override
+    async def unsubscribe(self) -> None:
+        """Unsubscribe from any data streams needed by this node."""
+        _ = await asyncio.gather(
+            self.left.unsubscribe(),
+            self.right.unsubscribe(),
+        )
+
 
 @dataclass(kw_only=True)
 class Sub(AstNode[QuantityT]):
@@ -261,12 +287,20 @@ def format(self, wrap: bool = False) -> str:
 
     @override
     async def subscribe(self) -> None:
-        """Subscribe to any data streams needed by the function."""
+        """Subscribe to any data streams needed by this node."""
         _ = await asyncio.gather(
             self.left.subscribe(),
             self.right.subscribe(),
         )
 
+    @override
+    async def unsubscribe(self) -> None:
+        """Unsubscribe from any data streams needed by this node."""
+        _ = await asyncio.gather(
+            self.left.unsubscribe(),
+            self.right.unsubscribe(),
+        )
+
 
 @dataclass(kw_only=True)
 class Mul(AstNode[QuantityT]):
@@ -325,12 +359,20 @@ def format(self, wrap: bool = False) -> str:
 
     @override
     async def subscribe(self) -> None:
-        """Subscribe to any data streams needed by the function."""
+        """Subscribe to any data streams needed by this node."""
         _ = await asyncio.gather(
             self.left.subscribe(),
             self.right.subscribe(),
         )
 
+    @override
+    async def unsubscribe(self) -> None:
+        """Unsubscribe from any data streams needed by this node."""
+        _ = await asyncio.gather(
+            self.left.unsubscribe(),
+            self.right.unsubscribe(),
+        )
+
 
 @dataclass(kw_only=True)
 class Div(AstNode[QuantityT]):
@@ -388,8 +430,16 @@ def format(self, wrap: bool = False) -> str:
 
     @override
     async def subscribe(self) -> None:
-        """Subscribe to any data streams needed by the function."""
+        """Subscribe to any data streams needed by this node."""
         _ = await asyncio.gather(
             self.left.subscribe(),
             self.right.subscribe(),
         )
+
+    @override
+    async def unsubscribe(self) -> None:
+        """Unsubscribe from any data streams needed by this node."""
+        _ = await asyncio.gather(
+            self.left.unsubscribe(),
+            self.right.unsubscribe(),
+        )

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

@@ -38,6 +38,10 @@ def __str__(self) -> str:
     async def subscribe(self) -> None:
         """Subscribe to any data streams needed by this node."""
 
+    @abc.abstractmethod
+    async def unsubscribe(self) -> None:
+        """Unsubscribe from any data streams used by this node."""
+
 
 class NodeSynchronizer(Generic[QuantityT]):
     """A helper class to synchronize multiple AST nodes."""

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

@@ -47,6 +47,12 @@ async def subscribe(self) -> None:
             *(param.subscribe() for param in self.params),
         )
 
+    async def unsubscribe(self) -> None:
+        """Unsubscribe from any data streams needed by the function."""
+        _ = await asyncio.gather(
+            *(param.unsubscribe() for param in self.params),
+        )
+
     @classmethod
     def from_string(
         cls, name: str, params: list[AstNode[QuantityT]]