Number cos/acos operations

finos · Jul 27, 2023 · 3a403c4 · 3a403c4
1 parent 9f6a76c
commit 3a403c4
Show file tree

Hide file tree

Showing 6 changed files with 118 additions and 0 deletions.
diff --git a/pylegend/core/databse/sql_to_string/db_extension.py b/pylegend/core/databse/sql_to_string/db_extension.py
@@ -94,6 +94,8 @@
     RoundExpression,
     SineExpression,
     ArcSineExpression,
+    CosineExpression,
+    ArcCosineExpression,
 )
 
 
@@ -372,6 +374,10 @@ def expression_processor(
         return extension.process_sine_expression(expression, config)
     elif isinstance(expression, ArcSineExpression):
         return extension.process_arc_sine_expression(expression, config)
+    elif isinstance(expression, CosineExpression):
+        return extension.process_cosine_expression(expression, config)
+    elif isinstance(expression, ArcCosineExpression):
+        return extension.process_arc_cosine_expression(expression, config)
     else:
         raise ValueError("Unsupported expression type: " + str(type(expression)))  # pragma: no cover
 
@@ -1048,6 +1054,16 @@ def process_arc_sine_expression(self, expr: ArcSineExpression, config: SqlToStri
             value=self.process_expression(expr.value, config)
         )
 
+    def process_cosine_expression(self, expr: CosineExpression, config: SqlToStringConfig) -> str:
+        return "COS({value})".format(
+            value=self.process_expression(expr.value, config)
+        )
+
+    def process_arc_cosine_expression(self, expr: ArcCosineExpression, config: SqlToStringConfig) -> str:
+        return "ACOS({value})".format(
+            value=self.process_expression(expr.value, config)
+        )
+
     def process_qualified_name(self, qualified_name: QualifiedName, config: SqlToStringConfig) -> str:
         return qualified_name_processor(qualified_name, self, config)
 

diff --git a/pylegend/core/language/operations/number_operation_expressions.py b/pylegend/core/language/operations/number_operation_expressions.py
@@ -46,6 +46,8 @@
     RoundExpression,
     SineExpression,
     ArcSineExpression,
+    CosineExpression,
+    ArcCosineExpression,
 )
 
 
@@ -71,6 +73,8 @@
     "PyLegendNumberRoundExpression",
     "PyLegendNumberSineExpression",
     "PyLegendNumberArcSineExpression",
+    "PyLegendNumberCosineExpression",
+    "PyLegendNumberArcCosineExpression",
 ]
 
 
@@ -493,3 +497,41 @@ def __init__(self, operand: PyLegendExpressionNumberReturn) -> None:
             operand,
             PyLegendNumberArcSineExpression.__to_sql_func
         )
+
+
+class PyLegendNumberCosineExpression(PyLegendUnaryExpression, PyLegendExpressionNumberReturn):
+
+    @staticmethod
+    def __to_sql_func(
+            expression: Expression,
+            frame_name_to_base_query_map: PyLegendDict[str, QuerySpecification],
+            config: FrameToSqlConfig
+    ) -> Expression:
+        return CosineExpression(expression)
+
+    def __init__(self, operand: PyLegendExpressionNumberReturn) -> None:
+        PyLegendExpressionNumberReturn.__init__(self)
+        PyLegendUnaryExpression.__init__(
+            self,
+            operand,
+            PyLegendNumberCosineExpression.__to_sql_func
+        )
+
+
+class PyLegendNumberArcCosineExpression(PyLegendUnaryExpression, PyLegendExpressionNumberReturn):
+
+    @staticmethod
+    def __to_sql_func(
+            expression: Expression,
+            frame_name_to_base_query_map: PyLegendDict[str, QuerySpecification],
+            config: FrameToSqlConfig
+    ) -> Expression:
+        return ArcCosineExpression(expression)
+
+    def __init__(self, operand: PyLegendExpressionNumberReturn) -> None:
+        PyLegendExpressionNumberReturn.__init__(self)
+        PyLegendUnaryExpression.__init__(
+            self,
+            operand,
+            PyLegendNumberArcCosineExpression.__to_sql_func
+        )
diff --git a/pylegend/core/language/primitives/number.py b/pylegend/core/language/primitives/number.py
@@ -50,6 +50,8 @@
     PyLegendNumberRoundExpression,
     PyLegendNumberSineExpression,
     PyLegendNumberArcSineExpression,
+    PyLegendNumberCosineExpression,
+    PyLegendNumberArcCosineExpression,
 )
 from pylegend.core.sql.metamodel import (
     Expression,
@@ -243,6 +245,12 @@ def sin(self) -> "PyLegendNumber":
     def asin(self) -> "PyLegendNumber":
         return PyLegendNumber(PyLegendNumberArcSineExpression(self.__value))
 
+    def cos(self) -> "PyLegendNumber":
+        return PyLegendNumber(PyLegendNumberCosineExpression(self.__value))
+
+    def acos(self) -> "PyLegendNumber":
+        return PyLegendNumber(PyLegendNumberArcCosineExpression(self.__value))
+
     def round(
             self,
             n: PyLegendOptional[int] = None

diff --git a/pylegend/core/sql/metamodel_extension.py b/pylegend/core/sql/metamodel_extension.py
@@ -42,6 +42,8 @@
     "RoundExpression",
     "SineExpression",
     "ArcSineExpression",
+    "CosineExpression",
+    "ArcCosineExpression",
 ]
 
 
@@ -279,3 +281,25 @@ def __init__(
     ) -> None:
         super().__init__(_type="arcSineExpression")
         self.value = value
+
+
+class CosineExpression(Expression):
+    value: "Expression"
+
+    def __init__(
+        self,
+        value: "Expression",
+    ) -> None:
+        super().__init__(_type="cosineExpression")
+        self.value = value
+
+
+class ArcCosineExpression(Expression):
+    value: "Expression"
+
+    def __init__(
+        self,
+        value: "Expression",
+    ) -> None:
+        super().__init__(_type="arcCosineExpression")
+        self.value = value
diff --git a/pylegend/tests/core/database/test_sql_to_string.py b/pylegend/tests/core/database/test_sql_to_string.py
@@ -91,6 +91,8 @@
     RoundExpression,
     SineExpression,
     ArcSineExpression,
+    CosineExpression,
+    ArcCosineExpression,
 )
 
 
@@ -1319,3 +1321,17 @@ def test_process_arc_sine_expression(self) -> None:
 
         expr = ArcSineExpression(IntegerLiteral(10))
         assert extension.process_expression(expr, config) == "ASIN(10)"
+
+    def test_process_cosine_expression(self) -> None:
+        extension = SqlToStringDbExtension()
+        config = SqlToStringConfig(SqlToStringFormat(pretty=False))
+
+        expr = CosineExpression(IntegerLiteral(10))
+        assert extension.process_expression(expr, config) == "COS(10)"
+
+    def test_process_arc_cosine_expression(self) -> None:
+        extension = SqlToStringDbExtension()
+        config = SqlToStringConfig(SqlToStringFormat(pretty=False))
+
+        expr = ArcCosineExpression(IntegerLiteral(10))
+        assert extension.process_expression(expr, config) == "ACOS(10)"
diff --git a/pylegend/tests/core/language/primitives/test_number.py b/pylegend/tests/core/language/primitives/test_number.py
@@ -209,6 +209,18 @@ def test_number_arc_sine_expr(self) -> None:
         assert self.__generate_sql_string(lambda x: (x.get_number("col2") + x.get_number("col1")).asin()) == \
                'ASIN(("root".col2 + "root".col1))'
 
+    def test_number_cosine_expr(self) -> None:
+        assert self.__generate_sql_string(lambda x: x.get_number("col2").cos()) == \
+               'COS("root".col2)'
+        assert self.__generate_sql_string(lambda x: (x.get_number("col2") + x.get_number("col1")).cos()) == \
+               'COS(("root".col2 + "root".col1))'
+
+    def test_number_arc_cosine_expr(self) -> None:
+        assert self.__generate_sql_string(lambda x: x.get_number("col2").acos()) == \
+               'ACOS("root".col2)'
+        assert self.__generate_sql_string(lambda x: (x.get_number("col2") + x.get_number("col1")).acos()) == \
+               'ACOS(("root".col2 + "root".col1))'
+
     def __generate_sql_string(self, f) -> str:  # type: ignore
         return self.db_extension.process_expression(
             f(self.tds_row).to_sql_expression({"t": self.base_query}, self.frame_to_sql_config),