pylegend 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

pylegend/core/tds/pandas_api/frames/functions/assign_function.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from datetime import date, datetime
+
 from pylegend._typing import (
     PyLegendList,
     PyLegendSequence,
@@ -20,17 +21,7 @@ from pylegend._typing import (
     PyLegendCallable,
     PyLegendUnion,
 )
-from pylegend.core.tds.pandas_api.frames.pandas_api_applied_function_tds_frame import PandasApiAppliedFunction
-from pylegend.core.tds.sql_query_helpers import copy_query, create_sub_query
-from pylegend.core.sql.metamodel import (
-    QuerySpecification,
-    SingleColumn,
-)
-from pylegend.core.tds.pandas_api.frames.pandas_api_base_tds_frame import PandasApiBaseTdsFrame
-from pylegend.core.tds.tds_column import TdsColumn, PrimitiveTdsColumn
-from pylegend.core.tds.tds_frame import FrameToSqlConfig
 from pylegend.core.language import (
-    LegacyApiTdsRow,
     PyLegendPrimitive,
     PyLegendInteger,
     PyLegendFloat,
@@ -38,13 +29,23 @@ from pylegend.core.language import (
     PyLegendBoolean,
     PyLegendString,
 )
+from pylegend.core.language.pandas_api.pandas_api_tds_row import PandasApiTdsRow
+from pylegend.core.sql.metamodel import (
+    QuerySpecification,
+    SingleColumn,
+)
+from pylegend.core.tds.pandas_api.frames.pandas_api_applied_function_tds_frame import PandasApiAppliedFunction
+from pylegend.core.tds.pandas_api.frames.pandas_api_base_tds_frame import PandasApiBaseTdsFrame
+from pylegend.core.tds.sql_query_helpers import copy_query, create_sub_query
+from pylegend.core.tds.tds_column import TdsColumn, PrimitiveTdsColumn
+from pylegend.core.tds.tds_frame import FrameToSqlConfig, FrameToPureConfig
 
 
 class AssignFunction(PandasApiAppliedFunction):
     __base_frame: PandasApiBaseTdsFrame
     __col_definitions: PyLegendDict[
         str,
-        PyLegendCallable[[LegacyApiTdsRow], PyLegendUnion[int, float, bool, str, date, datetime, PyLegendPrimitive]],
+        PyLegendCallable[[PandasApiTdsRow], PyLegendUnion[int, float, bool, str, date, datetime, PyLegendPrimitive]],
     ]
 
     @classmethod
@@ -56,7 +57,8 @@ class AssignFunction(PandasApiAppliedFunction):
             base_frame: PandasApiBaseTdsFrame,
             col_definitions: PyLegendDict[
                 str,
-                PyLegendCallable[[LegacyApiTdsRow], PyLegendUnion[int, float, bool, str, date, datetime, PyLegendPrimitive]],
+                PyLegendCallable[
+                    [PandasApiTdsRow], PyLegendUnion[int, float, bool, str, date, datetime, PyLegendPrimitive]],
             ]
     ) -> None:
         self.__base_frame = base_frame
@@ -72,7 +74,7 @@ class AssignFunction(PandasApiAppliedFunction):
             copy_query(base_query)
         )
 
-        tds_row = LegacyApiTdsRow.from_tds_frame("frame", self.__base_frame)
+        tds_row = PandasApiTdsRow.from_tds_frame("frame", self.__base_frame)
         for col, func in self.__col_definitions.items():
             res = func(tds_row)
             if not isinstance(res, PyLegendPrimitive):
@@ -86,6 +88,9 @@ class AssignFunction(PandasApiAppliedFunction):
             )
         return new_query
 
+    def to_pure(self, config: FrameToPureConfig) -> str:
+        raise NotImplementedError("to_pure is not implemented yet")  # pragma: no cover
+
     def base_frame(self) -> PandasApiBaseTdsFrame:
         return self.__base_frame
 
@@ -94,7 +99,7 @@ class AssignFunction(PandasApiAppliedFunction):
 
     def calculate_columns(self) -> PyLegendSequence["TdsColumn"]:
         new_cols = [c.copy() for c in self.__base_frame.columns()]
-        tds_row = LegacyApiTdsRow.from_tds_frame("frame", self.__base_frame)
+        tds_row = PandasApiTdsRow.from_tds_frame("frame", self.__base_frame)
         for col, func in self.__col_definitions.items():
             res = func(tds_row)
             if isinstance(res, (int, PyLegendInteger)):
@@ -112,7 +117,7 @@ class AssignFunction(PandasApiAppliedFunction):
         return new_cols
 
     def validate(self) -> bool:
-        tds_row = LegacyApiTdsRow.from_tds_frame("frame", self.__base_frame)
+        tds_row = PandasApiTdsRow.from_tds_frame("frame", self.__base_frame)
         for col, f in self.__col_definitions.items():
             f(tds_row)
         return True

pylegend/core/tds/pandas_api/frames/functions/drop.py

@@ -0,0 +1,171 @@
+# Copyright 2025 Goldman Sachs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pylegend._typing import (
+    PyLegendList,
+    PyLegendSet,
+    PyLegendSequence,
+    PyLegendUnion,
+    PyLegendOptional
+)
+from pylegend.core.language.shared.primitives.boolean import PyLegendBoolean
+from pylegend.core.language.shared.primitives.integer import PyLegendInteger
+from pylegend.core.sql.metamodel import (
+    QuerySpecification
+)
+from pylegend.core.tds.pandas_api.frames.functions.filter import PandasApiFilterFunction
+from pylegend.core.tds.pandas_api.frames.pandas_api_applied_function_tds_frame import PandasApiAppliedFunction
+from pylegend.core.tds.pandas_api.frames.pandas_api_base_tds_frame import PandasApiBaseTdsFrame
+from pylegend.core.tds.tds_column import TdsColumn
+from pylegend.core.tds.tds_frame import FrameToPureConfig, FrameToSqlConfig
+
+__all__: PyLegendSequence[str] = [
+    "PandasApiDropFunction"
+]
+
+
+class PandasApiDropFunction(PandasApiAppliedFunction):
+    __base_frame: PandasApiBaseTdsFrame
+    __labels: PyLegendOptional[PyLegendUnion[str, PyLegendSequence[str], PyLegendSet[str]]]
+    __axis: PyLegendUnion[str, int, PyLegendInteger]
+    __index: PyLegendOptional[PyLegendUnion[str, PyLegendSequence[str], PyLegendSet[str]]]
+    __columns: PyLegendOptional[PyLegendUnion[str, PyLegendSequence[str], PyLegendSet[str]]]
+    __level: PyLegendOptional[PyLegendUnion[int, PyLegendInteger, str]]
+    __inplace: PyLegendUnion[bool, PyLegendBoolean]
+    __errors: str
+
+    @classmethod
+    def name(cls) -> str:
+        return "drop"  # pragma: no cover
+
+    def __init__(
+            self,
+            base_frame: PandasApiBaseTdsFrame,
+            labels: PyLegendOptional[PyLegendUnion[str, PyLegendSequence[str], PyLegendSet[str]]],
+            axis: PyLegendUnion[str, int, PyLegendInteger],
+            index: PyLegendOptional[PyLegendUnion[str, PyLegendSequence[str], PyLegendSet[str]]],
+            columns: PyLegendOptional[PyLegendUnion[str, PyLegendSequence[str], PyLegendSet[str]]],
+            level: PyLegendOptional[PyLegendUnion[int, PyLegendInteger, str]],
+            inplace: PyLegendUnion[bool, PyLegendBoolean],
+            errors: str
+    ) -> None:
+        self.__base_frame = base_frame
+        self.__labels = labels
+        self.__axis = axis
+        self.__index = index
+        self.__columns = columns
+        self.__level = level
+        self.__inplace = inplace
+        self.__errors = errors
+
+    def to_sql(self, config: FrameToSqlConfig) -> QuerySpecification:
+        base_cols = [c.get_name() for c in self.__base_frame.columns()]
+
+        if self.__errors == "raise":
+            not_found = [col for col in self.__columns if col not in base_cols]  # type: ignore
+            if not_found:
+                raise KeyError(f"{not_found} not found in axis")
+
+        columns_to_retain = [col for col in base_cols if col not in self.__columns]  # type: ignore
+        filter_func = PandasApiFilterFunction(
+            base_frame=self.__base_frame,
+            items=columns_to_retain,
+            like=None,
+            regex=None,
+            axis=1
+        )
+        return filter_func.to_sql(config)
+
+    def to_pure(self, config: FrameToPureConfig) -> str:
+        base_cols = [c.get_name() for c in self.__base_frame.columns()]
+        if self.__errors == "raise":
+            not_found = [col for col in self.__columns if col not in base_cols]  # type: ignore
+            if not_found:
+                raise KeyError(f"{not_found} not found in axis")
+
+        columns_to_retain = [col for col in base_cols if col not in self.__columns]  # type: ignore
+        filter_func = PandasApiFilterFunction(
+            base_frame=self.__base_frame,
+            items=columns_to_retain,
+            like=None,
+            regex=None,
+            axis=1
+        )
+        return filter_func.to_pure(config)
+
+    def base_frame(self) -> PandasApiBaseTdsFrame:
+        return self.__base_frame
+
+    def tds_frame_parameters(self) -> PyLegendList["PandasApiBaseTdsFrame"]:
+        return []
+
+    def calculate_columns(self) -> PyLegendSequence["TdsColumn"]:
+        base_cols = [c.copy() for c in self.__base_frame.columns()]
+        if self.__columns is not None:
+            new_cols = []
+            for col in base_cols:
+                if col.get_name() not in self.__columns:
+                    new_cols.append(col.copy())
+            return new_cols
+        return base_cols  # pragma: no cover
+
+    def validate(self) -> bool:
+        valid_paramters: int = 0
+        if self.__axis is not None:
+            if isinstance(self.__axis, (str, int, PyLegendInteger)):
+                if self.__axis != 1 and self.__axis != "columns":
+                    if self.__axis == 0 or self.__axis == "index":
+                        raise NotImplementedError(
+                            f"Axis {self.__axis} is not supported for 'drop' function in PandasApi")
+                    else:
+                        raise ValueError(f"No axis named {self.__axis} for object type Tds DataFrame")
+            else:
+                raise TypeError(f"No axis named {self.__axis} for object type Tds DataFrame")  # pragma: no cover
+        if self.__level is not None:
+            raise NotImplementedError("'level' parameter is not supported for 'drop' function in PandasApi")
+
+        if self.__index is not None:
+            raise NotImplementedError("'index' parameter is not supported for 'drop' function in PandasApi")
+
+        if self.__labels is not None:
+            valid_paramters += 1
+
+            if self.__columns is None:
+                self.__columns = self.__labels
+            else:
+                raise ValueError("Cannot specify both 'labels' and 'columns'")
+
+        if self.__columns is not None:
+            def _normalize_columns(columns):  # type: ignore
+                if columns is None:
+                    return []  # pragma: no cover
+                if isinstance(columns, str):
+                    return [columns]
+                if isinstance(columns, (PyLegendSequence, PyLegendSet)):
+                    return list(columns)
+                raise TypeError(f"Unsupported type for columns: {type(columns)}")
+
+            valid_paramters += 1
+            self.__columns = _normalize_columns(self.__columns)  # type: ignore
+
+        if isinstance(self.__inplace, (bool, PyLegendBoolean)):
+            if self.__inplace is False:
+                raise NotImplementedError(f"Only inplace=True is supported. Got inplace={self.__inplace!r}")
+        else:
+            raise TypeError(f"Inplace must be True. Got inplace={self.__inplace!r}")  # pragma: no cover
+
+        if valid_paramters == 0:
+            raise ValueError("Need to specify at least one of 'labels' or 'columns'")
+
+        return True

pylegend/core/tds/pandas_api/frames/functions/filter.py

@@ -0,0 +1,193 @@
+# Copyright 2025 Goldman Sachs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+from pylegend._typing import (
+    PyLegendUnion,
+    PyLegendOptional,
+    PyLegendSequence,
+    PyLegendList,
+    PyLegendTuple
+)
+from pylegend.core.language import (
+    PyLegendInteger,
+)
+from pylegend.core.language.shared.helpers import escape_column_name
+from pylegend.core.sql.metamodel import (
+    QuerySpecification,
+    SingleColumn,
+    SelectItem
+)
+from pylegend.core.tds.pandas_api.frames.pandas_api_applied_function_tds_frame import (
+    PandasApiAppliedFunction,
+)
+from pylegend.core.tds.pandas_api.frames.pandas_api_base_tds_frame import (
+    PandasApiBaseTdsFrame,
+)
+from pylegend.core.tds.sql_query_helpers import copy_query, create_sub_query
+from pylegend.core.tds.tds_column import TdsColumn
+from pylegend.core.tds.tds_frame import FrameToSqlConfig, FrameToPureConfig
+
+__all__: PyLegendSequence[str] = ["PandasApiFilterFunction"]
+
+
+class PandasApiFilterFunction(PandasApiAppliedFunction):
+    __base_frame: PandasApiBaseTdsFrame
+    __items: PyLegendOptional[PyLegendList[str]]
+    __like: PyLegendOptional[str]
+    __regex: PyLegendOptional[str]
+    __axis: PyLegendOptional[PyLegendUnion[str, int, PyLegendInteger]]
+
+    @classmethod
+    def name(cls) -> str:
+        return "filter"  # pragma: no cover
+
+    def __init__(
+            self,
+            base_frame: PandasApiBaseTdsFrame,
+            items: PyLegendOptional[PyLegendList[str]],
+            like: PyLegendOptional[str],
+            regex: PyLegendOptional[str],
+            axis: PyLegendOptional[PyLegendUnion[str, int, PyLegendInteger]],
+    ) -> None:
+        self.__base_frame = base_frame
+        self.__items = items
+        self.__like = like
+        self.__regex = regex
+        self.__axis = 1 if axis is None else axis
+
+    def __get_desired_columns(
+            self, col_names: PyLegendSequence[str]
+    ) -> PyLegendSequence[str]:
+        if self.__items is not None:
+            return self.__items
+        elif self.__like is not None:
+            return [col for col in col_names if self.__like in col]
+        elif self.__regex is not None:
+            regex_pattern = re.compile(self.__regex)
+            return [col for col in col_names if regex_pattern.search(col)]
+
+        return []  # pragma: no cover
+
+    def to_sql(self, config: FrameToSqlConfig) -> QuerySpecification:
+        base_query = self.__base_frame.to_sql_query_object(config)
+        db_extension = config.sql_to_string_generator().get_db_extension()
+        columns_to_retain = [db_extension.quote_identifier(x) for x in
+                             self.__get_desired_columns([c.get_name() for c in self.__base_frame.columns()])]
+
+        sub_query_required = (
+                len(base_query.groupBy) > 0 or
+                len(base_query.orderBy) > 0 or
+                base_query.having is not None or
+                base_query.select.distinct
+        )
+
+        if sub_query_required:
+            new_query = create_sub_query(base_query, config, "root", columns_to_retain=columns_to_retain)
+            return new_query
+        else:
+            new_cols_with_index: PyLegendList[PyLegendTuple[int, SelectItem]] = []
+            for col in base_query.select.selectItems:
+                if not isinstance(col, SingleColumn):
+                    raise ValueError("Select operation not supported for queries with columns other than SingleColumn")
+                if col.alias is None:
+                    raise ValueError("Select operation not supported for queries with SingleColumns with missing alias")
+                if col.alias in columns_to_retain:
+                    new_cols_with_index.append((columns_to_retain.index(col.alias), col))
+
+            new_select_items = [y[1] for y in sorted(new_cols_with_index, key=lambda x: x[0])]
+            new_query = copy_query(base_query)
+            new_query.select.selectItems = new_select_items
+            return new_query
+
+    def to_pure(self, config: FrameToPureConfig) -> str:
+        col_names = [c.get_name() for c in self.__base_frame.columns()]
+        desired_columns = self.__get_desired_columns(col_names)
+        escaped_columns = [escape_column_name(col_name) for col_name in desired_columns]
+        return (
+            f"{self.__base_frame.to_pure(config)}{config.separator(1)}"
+            f"->select(~[{', '.join(escaped_columns)}])"
+        )
+
+    def base_frame(self) -> PandasApiBaseTdsFrame:
+        return self.__base_frame
+
+    def tds_frame_parameters(self) -> PyLegendList["PandasApiBaseTdsFrame"]:
+        return []
+
+    def calculate_columns(self) -> PyLegendSequence["TdsColumn"]:
+        base_cols = [c.copy() for c in self.__base_frame.columns()]
+        desired_col_names = self.__get_desired_columns([c.get_name() for c in base_cols])
+        return [
+            base_col.copy()
+            for base_col in base_cols
+            if base_col.get_name() in desired_col_names
+        ]
+
+    def validate(self) -> bool:
+        mutual_exclusion = sum(
+            [
+                self.__items is not None,
+                self.__like is not None,
+                self.__regex is not None,
+            ]
+        )
+        if mutual_exclusion > 1:
+            raise TypeError(
+                "Keyword arguments `items`, `like`, or `regex` are mutually exclusive"
+            )
+        if mutual_exclusion == 0:
+            raise TypeError("Must pass either `items`, `like`, or `regex`")
+
+        base_cols = [c.get_name() for c in self.__base_frame.columns()]
+        if self.__items is not None:
+            if not isinstance(self.__items, (list, PyLegendList)):
+                raise TypeError(
+                    f"Index(...) must be called with a collection, got '{self.__items}'"
+                )
+            invalid_cols = [item for item in self.__items if item not in base_cols]
+            if invalid_cols:
+                raise ValueError(
+                    f"Columns {invalid_cols} in `filter` items list do not exist. Available: {base_cols}"
+                )
+
+        if self.__like is not None:
+            if not isinstance(self.__like, str):
+                raise TypeError(f"'like' must be a string, got {type(self.__like)}")
+            if not any(self.__like in col for col in base_cols):
+                raise ValueError(
+                    f"No columns match the pattern '{self.__like}'. Available: {base_cols}"
+                )
+
+        if self.__regex is not None:
+            if not isinstance(self.__regex, str):
+                raise TypeError(f"'regex' must be a string, got {type(self.__regex)}")
+            try:
+                regex_pattern = re.compile(self.__regex)
+            except re.error as e:
+                raise ValueError(f"Invalid regex pattern '{self.__regex}': {e}")
+            if not any(regex_pattern.search(col) for col in base_cols):
+                raise ValueError(
+                    f"No columns match the regex '{self.__regex}'. Available: {base_cols}"
+                )
+
+        if not isinstance(
+                self.__axis, (str, int, PyLegendInteger)
+        ) or self.__axis not in [1, "columns"]:
+            raise ValueError(
+                f"Unsupported axis value: {self.__axis}. Expected 1 or 'columns'"
+            )
+
+        return True

pylegend/core/tds/pandas_api/frames/functions/filtering.py

@@ -0,0 +1,85 @@
+# Copyright 2025 Goldman Sachs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pylegend._typing import (
+    PyLegendSequence,
+    PyLegendList
+)
+from pylegend.core.language import PyLegendBoolean
+from pylegend.core.sql.metamodel import (
+    QuerySpecification,
+    LogicalBinaryExpression,
+    LogicalBinaryType,
+)
+from pylegend.core.tds.pandas_api.frames.pandas_api_applied_function_tds_frame import PandasApiAppliedFunction
+from pylegend.core.tds.pandas_api.frames.pandas_api_base_tds_frame import PandasApiBaseTdsFrame
+from pylegend.core.tds.sql_query_helpers import copy_query, create_sub_query
+from pylegend.core.tds.tds_column import TdsColumn
+from pylegend.core.tds.tds_frame import FrameToSqlConfig, FrameToPureConfig
+
+__all__: PyLegendSequence[str] = ["PandasApiFilteringFunction"]
+
+
+class PandasApiFilteringFunction(PandasApiAppliedFunction):
+    __base_frame: PandasApiBaseTdsFrame
+    __filter_expr: PyLegendBoolean
+
+    @classmethod
+    def name(cls) -> str:
+        return "boolean_filter"  # pragma: no cover
+
+    def __init__(
+            self,
+            base_frame: PandasApiBaseTdsFrame,
+            filter_expr: PyLegendBoolean
+    ) -> None:
+        self.__base_frame = base_frame
+        self.__filter_expr = filter_expr
+
+    def to_sql(self, config: FrameToSqlConfig) -> QuerySpecification:
+        base_query = self.__base_frame.to_sql_query_object(config)
+        should_create_sub_query = (len(base_query.groupBy) > 0) or \
+                                  (base_query.offset is not None) or (base_query.limit is not None)
+        new_query = (
+            create_sub_query(base_query, config, "root") if should_create_sub_query else
+            copy_query(base_query)
+        )
+
+        sql_expr = self.__filter_expr.to_sql_expression({"c": new_query}, config)
+
+        if new_query.where is None:
+            new_query.where = sql_expr
+        else:
+            new_query.where = LogicalBinaryExpression(
+                type_=LogicalBinaryType.AND,
+                left=new_query.where,
+                right=sql_expr
+            )
+        return new_query
+
+    def to_pure(self, config: FrameToPureConfig) -> str:
+        pure_expr = self.__filter_expr.to_pure_expression(config)
+        return f"{self.__base_frame.to_pure(config)}{config.separator(1)}->filter(c|{pure_expr})"
+
+    def base_frame(self) -> PandasApiBaseTdsFrame:
+        return self.__base_frame
+
+    def tds_frame_parameters(self) -> PyLegendList["PandasApiBaseTdsFrame"]:
+        return []
+
+    def calculate_columns(self) -> PyLegendSequence["TdsColumn"]:
+        return [c.copy() for c in self.__base_frame.columns()]
+
+    def validate(self) -> bool:
+        return True