Improve typing for HigherOrderFormulaBuilder by adding overloads

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

Author: shsms

src/frequenz/sdk/timeseries/_formula_engine/_formula_engine.py

@@ -547,6 +547,18 @@ def __init__(self, engine: FormulaEngine | FormulaEngine3Phase) -> None:
         ] = deque()
         self._steps.append((TokenType.COMPONENT_METRIC, engine))
 
+    @overload
+    def _push(
+        self, oper: str, other: _CompositionType1Phase
+    ) -> HigherOrderFormulaBuilder:
+        ...
+
+    @overload
+    def _push(
+        self, oper: str, other: _CompositionType3Phase
+    ) -> HigherOrderFormulaBuilder3Phase:
+        ...
+
     def _push(
         self, oper: str, other: _CompositionType
     ) -> HigherOrderFormulaBuilder | HigherOrderFormulaBuilder3Phase:
@@ -569,6 +581,14 @@ def _push(
         )
         return self
 
+    @overload
+    def __add__(self, other: _CompositionType1Phase) -> HigherOrderFormulaBuilder:
+        ...
+
+    @overload
+    def __add__(self, other: _CompositionType3Phase) -> HigherOrderFormulaBuilder3Phase:
+        ...
+
     def __add__(
         self, other: _CompositionType
     ) -> HigherOrderFormulaBuilder | HigherOrderFormulaBuilder3Phase:
@@ -584,6 +604,14 @@ def __add__(
         """
         return self._push("+", other)
 
+    @overload
+    def __sub__(self, other: _CompositionType1Phase) -> HigherOrderFormulaBuilder:
+        ...
+
+    @overload
+    def __sub__(self, other: _CompositionType3Phase) -> HigherOrderFormulaBuilder3Phase:
+        ...
+
     def __sub__(
         self,
         other: _CompositionType,
@@ -600,6 +628,14 @@ def __sub__(
         """
         return self._push("-", other)
 
+    @overload
+    def __mul__(self, other: _CompositionType1Phase) -> HigherOrderFormulaBuilder:
+        ...
+
+    @overload
+    def __mul__(self, other: _CompositionType3Phase) -> HigherOrderFormulaBuilder3Phase:
+        ...
+
     def __mul__(
         self,
         other: _CompositionType,
@@ -616,6 +652,16 @@ def __mul__(
         """
         return self._push("*", other)
 
+    @overload
+    def __truediv__(self, other: _CompositionType1Phase) -> HigherOrderFormulaBuilder:
+        ...
+
+    @overload
+    def __truediv__(
+        self, other: _CompositionType3Phase
+    ) -> HigherOrderFormulaBuilder3Phase:
+        ...
+
     def __truediv__(
         self,
         other: _CompositionType,