n-ary operations

pyproject.toml

@@ -1,7 +1,7 @@
 [tool]
 [tool.poetry]
 name = "ep-stats"
-version = "2.5.2"
+version = "2.5.3"
 homepage = "https://github.com/avast/ep-stats"
 description = "Statistical package to evaluate ab tests in experimentation platform."
 authors = [

src/epstats/toolkit/parser.py

@@ -1,4 +1,5 @@
 from collections import Counter
+from functools import reduce
 from typing import Set
 
 import pandas as pd
@@ -49,11 +50,11 @@ def __init__(self, nominator: str, denominator: str):
         expr = infixNotation(
             operand,
             [
-                (multop, 2, opAssoc.LEFT, MultBinOp),
-                (divop, 2, opAssoc.LEFT, DivBinOp),
-                (subop, 2, opAssoc.LEFT, SubBinOp),
-                (plusop, 2, opAssoc.LEFT, PlusBinOp),
-                (tildaop, 2, opAssoc.LEFT, TildaBinOp),
+                (multop, 2, opAssoc.LEFT, MultOp),
+                (divop, 2, opAssoc.LEFT, DivOp),
+                (subop, 2, opAssoc.LEFT, SubOp),
+                (plusop, 2, opAssoc.LEFT, PlusOp),
+                (tildaop, 2, opAssoc.LEFT, TildaOp),
             ],
         )
 
@@ -332,7 +333,7 @@ def is_dimensional(self):
     __repr__ = __str__
 
 
-class BinOp:
+class Op:
     """
     Operation connecting `EpGoal` or `Number` terms in nominator or denominator expression
     eg. `value(test_unit_type.unit.conversion) - value(test_unit_type.unit.refund)`.
@@ -344,14 +345,18 @@ def __init__(self, t):
     def symbol(self):
         raise NotImplementedError()
 
-    def evaluate_agg(self, goals):
+    @staticmethod
+    def reduce_f(x, y):
         raise NotImplementedError()
 
+    def evaluate_agg(self, goals):
+        return reduce(self.reduce_f, [arg.evaluate_agg(goals) for arg in self.args])
+
     def evaluate_by_unit(self, goals):
-        raise NotImplementedError()
+        return reduce(self.reduce_f, [arg.evaluate_by_unit(goals) for arg in self.args])
 
     def evaluate_sqr(self, goals):
-        raise NotImplementedError()
+        return reduce(self.reduce_f, [arg.evaluate_sqr(goals) for arg in self.args])
 
     def get_goals_str(self) -> Set[str]:
         return set().union(*map(lambda o: o.get_goals_str(), self.args))
@@ -366,76 +371,56 @@ def __str__(self):
     __repr__ = __str__
 
 
-class PlusBinOp(BinOp):
+class PlusOp(Op):
     def symbol(self):
         return "+"
 
-    def evaluate_agg(self, goals):
-        return self.args[0].evaluate_agg(goals) + self.args[1].evaluate_agg(goals)
-
-    def evaluate_sqr(self, goals):
-        return self.args[0].evaluate_sqr(goals) + self.args[1].evaluate_sqr(goals)
-
-    def evaluate_by_unit(self, goals):
-        return self.args[0].evaluate_by_unit(goals) + self.args[1].evaluate_by_unit(goals)
+    @staticmethod
+    def reduce_f(x, y):
+        return x + y
 
 
-class MultBinOp(BinOp):
+class MultOp(Op):
     def symbol(self):
         return "*"
 
-    def evaluate_agg(self, goals):
-        return self.args[0].evaluate_agg(goals) * self.args[1].evaluate_agg(goals)
-
-    def evaluate_sqr(self, goals):
-        return self.args[0].evaluate_sqr(goals) * self.args[1].evaluate_sqr(goals)
-
-    def evaluate_by_unit(self, goals):
-        return self.args[0].evaluate_by_unit(goals) * self.args[1].evaluate_by_unit(goals)
+    @staticmethod
+    def reduce_f(x, y):
+        return x * y
 
 
-class DivBinOp(BinOp):
+class DivOp(Op):
     def symbol(self):
         return "/"
 
-    def evaluate_agg(self, goals):
-        return self.args[0].evaluate_agg(goals) / self.args[1].evaluate_agg(goals)
-
-    def evaluate_sqr(self, goals):
-        return self.args[0].evaluate_sqr(goals) / self.args[1].evaluate_sqr(goals)
-
-    def evaluate_by_unit(self, goals):
-        return self.args[0].evaluate_by_unit(goals) / self.args[1].evaluate_by_unit(goals)
+    @staticmethod
+    def reduce_f(x, y):
+        return x / y
 
 
-class SubBinOp(BinOp):
+class SubOp(Op):
     def symbol(self):
         return "-"
 
-    def evaluate_agg(self, goals):
-        return self.args[0].evaluate_agg(goals) - self.args[1].evaluate_agg(goals)
-
-    def evaluate_sqr(self, goals):
-        return self.args[0].evaluate_sqr(goals) - self.args[1].evaluate_sqr(goals)
-
-    def evaluate_by_unit(self, goals):
-        return self.args[0].evaluate_by_unit(goals) - self.args[1].evaluate_by_unit(goals)
+    @staticmethod
+    def reduce_f(x, y):
+        return x - y
 
 
-class TildaBinOp(BinOp):
+class TildaOp(Op):
     """
-    Tilda treats the second operand as negative,
+    Tilda treats the following operands as negative,
     resulting in substraction of values and addition of squared values.
     """
 
     def symbol(self):
         return "~"
 
     def evaluate_agg(self, goals):
-        return self.args[0].evaluate_agg(goals) - self.args[1].evaluate_agg(goals)
+        return reduce(lambda x, y: x - y, [arg.evaluate_agg(goals) for arg in self.args])
 
     def evaluate_sqr(self, goals):
-        return self.args[0].evaluate_sqr(goals) + self.args[1].evaluate_sqr(goals)
+        return reduce(lambda x, y: x + y, [arg.evaluate_sqr(goals) for arg in self.args])
 
     def evaluate_by_unit(self, goals):
-        return self.args[0].evaluate_by_unit(goals) - self.args[1].evaluate_by_unit(goals)
+        return reduce(lambda x, y: x - y, [arg.evaluate_by_unit(goals) for arg in self.args])

tests/epstats/toolkit/test_parser.py

@@ -4,7 +4,7 @@
 from numpy.testing import assert_almost_equal
 from pyparsing import ParseException
 
-from src.epstats.toolkit.parser import MultBinOp, Parser
+from src.epstats.toolkit.parser import MultOp, Parser
 
 
 def test_evaluate_agg():
@@ -75,6 +75,21 @@ def test_evaluate_agg():
         conversion_sqr_value,
     )
 
+    parser = Parser(
+        """
+        value(test_unit_type.unit.conversion)
+        - value(test_unit_type.unit.refund)
+        - value(test_unit_type.unit.refund)
+        """,
+        "count(test_unit_type.unit.exposure)",
+    )
+    assert_count_value(
+        parser.evaluate_agg(goals),
+        goals[goals.goal == "exposure"]["count"],
+        conversion_value - refund_value - refund_value,
+        conversion_sqr_value - refund_sqr_value - refund_sqr_value,
+    )
+
 
 def test_evaluate_agg_dimensional():
     goals = pd.DataFrame(
@@ -357,7 +372,7 @@ def test_operator_position_not_correct(dimension_value):
 def test_numbers(nominator):
     assert isinstance(
         Parser(nominator, "count(test_unit_type.unit.conversion)")._nominator_expr,
-        MultBinOp,
+        MultOp,
     )