Extract Node base class to separate module

Also rename it to AstNode

This way, nodes can be accessed by the _function.py module, without
circular imports.  That would enable functions to handle their own
arguments.

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

Author: shsms

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

@@ -5,7 +5,6 @@
 
 from __future__ import annotations
 
-import abc
 import logging
 import math
 from collections.abc import AsyncIterator
@@ -16,33 +15,14 @@
 
 from ..._internal._math import is_close_to_zero
 from .._base_types import QuantityT, Sample
+from ._base_ast_node import AstNode
 from ._functions import Function
 
 _logger = logging.getLogger(__name__)
 
 
 @dataclass(kw_only=True)
-class Node(abc.ABC):
-    """An abstract syntax tree node representing a formula expression."""
-
-    span: tuple[int, int] | None = None
-
-    @abc.abstractmethod
-    def evaluate(self) -> float | None:
-        """Evaluate the expression and return its numerical value."""
-
-    @abc.abstractmethod
-    def format(self, wrap: bool = False) -> str:
-        """Return a string representation of the node."""
-
-    @override
-    def __str__(self) -> str:
-        """Return the string representation of the node."""
-        return self.format()
-
-
-@dataclass(kw_only=True)
-class TelemetryStream(Node, Generic[QuantityT]):
+class TelemetryStream(AstNode, Generic[QuantityT]):
     """A AST node that retrieves values from a component's telemetry stream."""
 
     source: str
@@ -74,11 +54,11 @@ async def fetch_next(self) -> None:
 
 
 @dataclass(kw_only=True)
-class FunCall(Node):
+class FunCall(AstNode):
     """A function call in the formula."""
 
     function: Function
-    args: list[Node]
+    args: list[AstNode]
 
     @override
     def evaluate(self) -> float | None:
@@ -93,7 +73,7 @@ def format(self, wrap: bool = False) -> str:
 
 
 @dataclass(kw_only=True)
-class Constant(Node):
+class Constant(AstNode):
     """A constant numerical value in the formula."""
 
     value: float
@@ -110,11 +90,11 @@ def format(self, wrap: bool = False) -> str:
 
 
 @dataclass(kw_only=True)
-class Add(Node):
+class Add(AstNode):
     """Addition operation node."""
 
-    left: Node
-    right: Node
+    left: AstNode
+    right: AstNode
 
     @override
     def evaluate(self) -> float | None:
@@ -135,11 +115,11 @@ def format(self, wrap: bool = False) -> str:
 
 
 @dataclass(kw_only=True)
-class Sub(Node):
+class Sub(AstNode):
     """Subtraction operation node."""
 
-    left: Node
-    right: Node
+    left: AstNode
+    right: AstNode
 
     @override
     def evaluate(self) -> float | None:
@@ -160,11 +140,11 @@ def format(self, wrap: bool = False) -> str:
 
 
 @dataclass(kw_only=True)
-class Mul(Node):
+class Mul(AstNode):
     """Multiplication operation node."""
 
-    left: Node
-    right: Node
+    left: AstNode
+    right: AstNode
 
     @override
     def evaluate(self) -> float | None:
@@ -182,11 +162,11 @@ def format(self, wrap: bool = False) -> str:
 
 
 @dataclass(kw_only=True)
-class Div(Node):
+class Div(AstNode):
     """Division operation node."""
 
-    left: Node
-    right: Node
+    left: AstNode
+    right: AstNode
 
     @override
     def evaluate(self) -> float | None:

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

@@ -0,0 +1,29 @@
+# License: MIT
+# Copyright © 2025 Frequenz Energy-as-a-Service GmbH
+
+"""Formula AST node base class."""
+
+import abc
+from dataclasses import dataclass
+
+from typing_extensions import override
+
+
+@dataclass(kw_only=True)
+class AstNode(abc.ABC):
+    """An abstract syntax tree node representing a formula expression."""
+
+    span: tuple[int, int] | None = None
+
+    @abc.abstractmethod
+    def evaluate(self) -> float | None:
+        """Evaluate the expression and return its numerical value."""
+
+    @abc.abstractmethod
+    def format(self, wrap: bool = False) -> str:
+        """Return a string representation of the node."""
+
+    @override
+    def __str__(self) -> str:
+        """Return the string representation of the node."""
+        return self.format()

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

@@ -20,6 +20,7 @@
 from .. import ReceiverFetcher, Sample
 from .._base_types import QuantityT
 from . import _ast
+from ._base_ast_node import AstNode
 from ._formula_evaluator import FormulaEvaluatingActor
 from ._functions import Coalesce, Max, Min
 
@@ -33,7 +34,7 @@ def __init__(  # pylint: disable=too-many-arguments
         self,
         *,
         name: str,
-        root: _ast.Node,
+        root: AstNode,
         create_method: Callable[[float], QuantityT],
         streams: list[_ast.TelemetryStream[QuantityT]],
         sub_formulas: list[Formula[QuantityT]] | None = None,
@@ -54,7 +55,7 @@ def __init__(  # pylint: disable=too-many-arguments
         """
         BackgroundService.__init__(self)
         self._name: str = name
-        self._root: _ast.Node = root
+        self._root: AstNode = root
         self._components: list[_ast.TelemetryStream[QuantityT]] = streams
         self._create_method: Callable[[float], QuantityT] = create_method
         self._sub_formulas: list[Formula[QuantityT]] = sub_formulas or []
@@ -154,7 +155,7 @@ class FormulaBuilder(Generic[QuantityT]):
 
     def __init__(
         self,
-        formula: Formula[QuantityT] | _ast.Node,
+        formula: Formula[QuantityT] | AstNode,
         create_method: Callable[[float], QuantityT],
         streams: list[_ast.TelemetryStream[QuantityT]] | None = None,
         sub_formulas: list[Formula[QuantityT]] | None = None,
@@ -176,7 +177,7 @@ def __init__(
         """Sub-formulas whose lifetimes are managed by this formula."""
 
         if isinstance(formula, Formula):
-            self.root: _ast.Node = _ast.TelemetryStream(
+            self.root: AstNode = _ast.TelemetryStream(
                 source=str(formula),
                 stream=formula.new_receiver(),
             )
@@ -266,7 +267,7 @@ def coalesce(
         other: list[FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT]],
     ) -> FormulaBuilder[QuantityT]:
         """Create a coalesce operation node."""
-        right_nodes: list[_ast.Node] = []
+        right_nodes: list[AstNode] = []
         for item in other:
             if isinstance(item, FormulaBuilder):
                 right_nodes.append(item.root)
@@ -299,7 +300,7 @@ def min(
         other: list[FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT]],
     ) -> FormulaBuilder[QuantityT]:
         """Create a min operation node."""
-        right_nodes: list[_ast.Node] = []
+        right_nodes: list[AstNode] = []
         for item in other:
             if isinstance(item, FormulaBuilder):
                 right_nodes.append(item.root)
@@ -332,7 +333,7 @@ def max(
         other: list[FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT]],
     ) -> FormulaBuilder[QuantityT]:
         """Create a max operation node."""
-        right_nodes: list[_ast.Node] = []
+        right_nodes: list[AstNode] = []
         for item in other:
             if isinstance(item, FormulaBuilder):
                 right_nodes.append(item.root)

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

@@ -16,6 +16,7 @@
 from ...actor import Actor
 from .._base_types import QuantityT, Sample
 from . import _ast
+from ._base_ast_node import AstNode
 from ._resampled_stream_fetcher import ResampledStreamFetcher
 
 _logger = logging.getLogger(__name__)
@@ -27,7 +28,7 @@ class FormulaEvaluatingActor(Generic[QuantityT], Actor):
     def __init__(  # pylint: disable=too-many-arguments
         self,
         *,
-        root: _ast.Node,
+        root: AstNode,
         components: list[_ast.TelemetryStream[QuantityT]],
         create_method: Callable[[float], QuantityT],
         output_channel: Broadcast[Sample[QuantityT]],
@@ -47,7 +48,7 @@ def __init__(  # pylint: disable=too-many-arguments
         """
         super().__init__()
 
-        self._root: _ast.Node = root
+        self._root: AstNode = root
         self._components: list[_ast.TelemetryStream[QuantityT]] = components
         self._create_method: Callable[[float], QuantityT] = create_method
         self._metric_fetcher: ResampledStreamFetcher | None = metric_fetcher

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

@@ -13,6 +13,7 @@
 from frequenz.sdk.timeseries._base_types import QuantityT
 
 from . import _ast, _token
+from ._base_ast_node import AstNode
 from ._formula import Formula
 from ._functions import Function
 from ._lexer import Lexer
@@ -63,7 +64,7 @@ def __init__(
         self._components: list[_ast.TelemetryStream[QuantityT]] = []
         self._create_method: Callable[[float], QuantityT] = create_method
 
-    def _parse_term(self) -> _ast.Node | None:
+    def _parse_term(self) -> AstNode | None:
         factor = self._parse_factor()
         if factor is None:
             return None
@@ -87,7 +88,7 @@ def _parse_term(self) -> _ast.Node | None:
 
         return factor
 
-    def _parse_factor(self) -> _ast.Node | None:
+    def _parse_factor(self) -> AstNode | None:
         unary = self._parse_unary()
 
         if unary is None:
@@ -109,11 +110,11 @@ def _parse_factor(self) -> _ast.Node | None:
 
         return unary
 
-    def _parse_unary(self) -> _ast.Node | None:
+    def _parse_unary(self) -> AstNode | None:
         token: _token.Token | None = self._lexer.peek()
         if token is not None and isinstance(token, _token.Minus):
             token = next(self._lexer)
-            primary: _ast.Node | None = self._parse_primary()
+            primary: AstNode | None = self._parse_primary()
             if primary is None:
                 raise ValueError(
                     f"Expected primary expression after unary '-' at position {token.span}"
@@ -124,11 +125,11 @@ def _parse_unary(self) -> _ast.Node | None:
 
         return self._parse_primary()
 
-    def _parse_bracketed(self) -> _ast.Node | None:
+    def _parse_bracketed(self) -> AstNode | None:
         oparen = next(self._lexer)  # consume '('
         assert isinstance(oparen, _token.OpenParen)
 
-        expr: _ast.Node | None = self._parse_term()
+        expr: AstNode | None = self._parse_term()
         if expr is None:
             raise ValueError(f"Expected expression after '(' at position {oparen.span}")
 
@@ -140,9 +141,9 @@ def _parse_bracketed(self) -> _ast.Node | None:
 
         return expr
 
-    def _parse_function_call(self) -> _ast.Node | None:
+    def _parse_function_call(self) -> AstNode | None:
         fn_name: _token.Token = next(self._lexer)
-        args: list[_ast.Node] = []
+        args: list[AstNode] = []
 
         token: _token.Token | None = self._lexer.peek()
         if token is None or not isinstance(token, _token.OpenParen):
@@ -176,7 +177,7 @@ def _parse_function_call(self) -> _ast.Node | None:
             args=args,
         )
 
-    def _parse_primary(self) -> _ast.Node | None:
+    def _parse_primary(self) -> AstNode | None:
         token: _token.Token | None = self._lexer.peek()
         if token is None:
             return None