diff --git a/pylegend/core/databse/sql_to_string/db_extension.py b/pylegend/core/databse/sql_to_string/db_extension.py index 86a8dfc8..1f6bafc1 100644 --- a/pylegend/core/databse/sql_to_string/db_extension.py +++ b/pylegend/core/databse/sql_to_string/db_extension.py @@ -98,6 +98,7 @@ ArcCosineExpression, TanExpression, ArcTanExpression, + CotExpression, ) @@ -384,6 +385,8 @@ def expression_processor( return extension.process_tan_expression(expression, config) elif isinstance(expression, ArcTanExpression): return extension.process_arc_tan_expression(expression, config) + elif isinstance(expression, CotExpression): + return extension.process_cot_expression(expression, config) else: raise ValueError("Unsupported expression type: " + str(type(expression))) # pragma: no cover @@ -1080,6 +1083,11 @@ def process_arc_tan_expression(self, expr: ArcTanExpression, config: SqlToString value=self.process_expression(expr.value, config) ) + def process_cot_expression(self, expr: CotExpression, config: SqlToStringConfig) -> str: + return "COT({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 index 24d97a8d..f7694c46 100644 --- a/pylegend/core/language/operations/number_operation_expressions.py +++ b/pylegend/core/language/operations/number_operation_expressions.py @@ -50,6 +50,7 @@ ArcCosineExpression, TanExpression, ArcTanExpression, + CotExpression, ) @@ -79,6 +80,7 @@ "PyLegendNumberArcCosineExpression", "PyLegendNumberTanExpression", "PyLegendNumberArcTanExpression", + "PyLegendNumberCotExpression", ] @@ -577,3 +579,22 @@ def __init__(self, operand: PyLegendExpressionNumberReturn) -> None: operand, PyLegendNumberArcTanExpression.__to_sql_func ) + + +class PyLegendNumberCotExpression(PyLegendUnaryExpression, PyLegendExpressionNumberReturn): + + @staticmethod + def __to_sql_func( + expression: Expression, + frame_name_to_base_query_map: PyLegendDict[str, QuerySpecification], + config: FrameToSqlConfig + ) -> Expression: + return CotExpression(expression) + + def __init__(self, operand: PyLegendExpressionNumberReturn) -> None: + PyLegendExpressionNumberReturn.__init__(self) + PyLegendUnaryExpression.__init__( + self, + operand, + PyLegendNumberCotExpression.__to_sql_func + ) diff --git a/pylegend/core/language/primitives/number.py b/pylegend/core/language/primitives/number.py index 7010de1c..660007d6 100644 --- a/pylegend/core/language/primitives/number.py +++ b/pylegend/core/language/primitives/number.py @@ -54,6 +54,7 @@ PyLegendNumberArcCosineExpression, PyLegendNumberTanExpression, PyLegendNumberArcTanExpression, + PyLegendNumberCotExpression, ) from pylegend.core.sql.metamodel import ( Expression, @@ -259,6 +260,9 @@ def tan(self) -> "PyLegendNumber": def atan(self) -> "PyLegendNumber": return PyLegendNumber(PyLegendNumberArcTanExpression(self.__value)) + def cot(self) -> "PyLegendNumber": + return PyLegendNumber(PyLegendNumberCotExpression(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 index 8e27c5f0..3322d506 100644 --- a/pylegend/core/sql/metamodel_extension.py +++ b/pylegend/core/sql/metamodel_extension.py @@ -46,6 +46,7 @@ "ArcCosineExpression", "TanExpression", "ArcTanExpression", + "CotExpression", ] @@ -327,3 +328,14 @@ def __init__( ) -> None: super().__init__(_type="arcTanExpression") self.value = value + + +class CotExpression(Expression): + value: "Expression" + + def __init__( + self, + value: "Expression", + ) -> None: + super().__init__(_type="cotExpression") + 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 index 4c941ade..45e6c6d9 100644 --- a/pylegend/tests/core/database/test_sql_to_string.py +++ b/pylegend/tests/core/database/test_sql_to_string.py @@ -95,6 +95,7 @@ ArcCosineExpression, TanExpression, ArcTanExpression, + CotExpression, ) @@ -1351,3 +1352,10 @@ def test_process_arc_tan_expression(self) -> None: expr = ArcTanExpression(DoubleLiteral(0.5)) assert extension.process_expression(expr, config) == "ATAN(0.5)" + + def test_process_cot_expression(self) -> None: + extension = SqlToStringDbExtension() + config = SqlToStringConfig(SqlToStringFormat(pretty=False)) + + expr = CotExpression(IntegerLiteral(10)) + assert extension.process_expression(expr, config) == "COT(10)" diff --git a/pylegend/tests/core/language/primitives/test_number.py b/pylegend/tests/core/language/primitives/test_number.py index 911c1a4a..88afce50 100644 --- a/pylegend/tests/core/language/primitives/test_number.py +++ b/pylegend/tests/core/language/primitives/test_number.py @@ -233,6 +233,12 @@ def test_number_arc_tan_expr(self) -> None: assert self.__generate_sql_string(lambda x: (x.get_number("col2") + x.get_number("col1")).atan()) == \ 'ATAN(("root".col2 + "root".col1))' + def test_number_cot_expr(self) -> None: + assert self.__generate_sql_string(lambda x: x.get_number("col2").cot()) == \ + 'COT("root".col2)' + assert self.__generate_sql_string(lambda x: (x.get_number("col2") + x.get_number("col1")).cot()) == \ + 'COT(("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),