pylegend 0.10.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

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

@@ -0,0 +1,214 @@
+# 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,
+    PyLegendSequence,
+    PyLegendUnion,
+    PyLegendOptional,
+    PyLegendCallable,
+    PyLegendDict
+)
+from pylegend.core.language import (
+    PyLegendInteger,
+    PyLegendBoolean
+)
+from pylegend.core.sql.metamodel import (
+    QuerySpecification,
+    SelectItem,
+    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
+from pylegend.core.tds.tds_column import TdsColumn
+from pylegend.core.tds.tds_frame import FrameToPureConfig, FrameToSqlConfig
+
+
+class PandasApiRenameFunction(PandasApiAppliedFunction):
+    __base_frame: PandasApiBaseTdsFrame
+    __mapper: PyLegendOptional[PyLegendUnion[PyLegendDict[str, str], PyLegendCallable[[str], str]]]
+    __axis: PyLegendUnion[str, int, PyLegendInteger]
+    __index: PyLegendOptional[PyLegendUnion[PyLegendDict[str, str], PyLegendCallable[[str], str]]]
+    __columns: PyLegendOptional[PyLegendUnion[PyLegendDict[str, str], PyLegendCallable[[str], str]]]
+    __level: PyLegendOptional[PyLegendUnion[int, PyLegendInteger, str]]
+    __inplace: PyLegendUnion[bool, PyLegendBoolean]
+    __copy: PyLegendUnion[bool, PyLegendBoolean]
+    __errors: str
+
+    @classmethod
+    def name(cls) -> str:
+        return "rename"  # pragma: no cover
+
+    def __init__(
+            self,
+            base_frame: PandasApiBaseTdsFrame,
+            mapper: PyLegendOptional[PyLegendUnion[PyLegendDict[str, str], PyLegendCallable[[str], str]]],
+            axis: PyLegendUnion[str, int, PyLegendInteger],
+            index: PyLegendOptional[PyLegendUnion[PyLegendDict[str, str], PyLegendCallable[[str], str]]],
+            columns: PyLegendOptional[PyLegendUnion[PyLegendDict[str, str], PyLegendCallable[[str], str]]],
+            level: PyLegendOptional[PyLegendUnion[int, PyLegendInteger, str]],
+            inplace: PyLegendUnion[bool, PyLegendBoolean],
+            errors: str,
+            copy: PyLegendUnion[bool, PyLegendBoolean]
+    ) -> None:
+        self.__base_frame = base_frame
+        self.__mapper = mapper
+        self.__axis = axis
+        self.__index = index
+        self.__columns = columns
+        self.__level = level
+        self.__inplace = inplace
+        self.__errors = errors
+        self.__copy = copy
+
+    def __resolve_columns_mapping(self) -> PyLegendDict[str, str]:
+        base_cols = [c.get_name() for c in self.__base_frame.columns()]
+        mapping_source = None
+
+        axis_is_columns = (self.__axis == 1 or self.__axis == "columns")
+
+        # Priority: explicit columns=, else mapper when axis targets columns
+        if self.__columns is not None:
+            mapping_source = self.__columns
+        elif self.__mapper is not None and axis_is_columns:
+            mapping_source = self.__mapper
+
+        if mapping_source is None:
+            return {}
+
+        if not callable(mapping_source) and not isinstance(mapping_source, dict):
+            raise TypeError(
+                f"Rename mapping must be a dict or a callable, got {type(mapping_source)}"
+            )
+
+        out: PyLegendDict[str, str] = {}
+        if callable(mapping_source):
+            func = mapping_source
+            for col in base_cols:
+                new = func(col)
+                if not isinstance(new, str):
+                    raise TypeError(
+                        f"Rename function must return str, got {type(new)} for column {col}")  # pragma: no cover
+                if new != col:
+                    out[col] = new
+        else:
+            # dict-like
+            dict_map: PyLegendDict[str, str] = mapping_source
+            if self.__errors == "raise":
+                missing = [k for k in dict_map.keys() if k not in base_cols]
+                if missing:
+                    raise KeyError(f"{missing} not found in axis")
+
+            for k, v in dict_map.items():
+                if k in base_cols and k != v:
+                    out[k] = v
+
+        return out
+
+    def to_sql(self, config: FrameToSqlConfig) -> QuerySpecification:
+        rename_map = self.__resolve_columns_mapping()
+        base_query = self.__base_frame.to_sql_query_object(config)
+        db_extension = config.sql_to_string_generator().get_db_extension()
+
+        # Prepare quoted lookup for aliases
+        quoted_from = [db_extension.quote_identifier(s) for s in rename_map.keys()]
+        quoted_to = [db_extension.quote_identifier(rename_map[s]) for s in rename_map.keys()]
+
+        new_select_items: PyLegendList[SelectItem] = []
+        for col in base_query.select.selectItems:
+            if not isinstance(col, SingleColumn):
+                raise ValueError("Rename operation not supported for non-SingleColumn select items")  # pragma: no cover
+            if col.alias is None:
+                raise ValueError("Rename operation requires SingleColumn items with aliases")  # pragma: no cover
+            if col.alias in quoted_from:
+                new_alias = quoted_to[quoted_from.index(col.alias)]
+                new_select_items.append(SingleColumn(alias=new_alias, expression=col.expression))
+            else:
+                new_select_items.append(col)
+
+        new_query = copy_query(base_query)
+        new_query.select.selectItems = new_select_items
+        return new_query
+
+    def to_pure(self, config: FrameToPureConfig) -> str:
+        rename_map = self.__resolve_columns_mapping()
+        base_pure = self.__base_frame.to_pure(config)
+
+        # Build a single project that aliases columns to new names
+        project_items: PyLegendList[str] = []
+        for c in self.__base_frame.columns():
+            orig = c.get_name()
+            new = rename_map.get(orig, orig)
+            project_items.append(f"{new}:x|$x.{orig}")
+
+        project_body = ", ".join(project_items)
+        return (
+            f"{base_pure}{config.separator(1)}"
+            f"->project({config.separator(2)}~[{project_body}]{config.separator(1)})"
+        )
+
+    def base_frame(self) -> PandasApiBaseTdsFrame:
+        return self.__base_frame
+
+    def tds_frame_parameters(self) -> PyLegendList["PandasApiBaseTdsFrame"]:
+        return []
+
+    def calculate_columns(self) -> PyLegendSequence["TdsColumn"]:
+        rename_map = self.__resolve_columns_mapping()
+        new_cols = []
+        for c in self.__base_frame.columns():
+            name = c.get_name()
+            if name in rename_map:
+                new_cols.append(c.copy_with_changed_name(rename_map[name]))
+            else:
+                new_cols.append(c.copy())
+        names = [c.get_name() for c in new_cols]
+        if len(names) != len(set(names)):
+            raise ValueError("Resulting columns contain duplicates after rename")
+        return new_cols
+
+    def validate(self) -> bool:
+        if self.__level is not None:
+            raise NotImplementedError("level parameter not supported yet in Pandas API")
+
+        if not isinstance(self.__inplace, bool):
+            raise TypeError(f"inplace must be bool. Got {type(self.__inplace)}")
+        if self.__inplace is True:
+            raise NotImplementedError("inplace=True not supported yet in Pandas API")
+
+        if not isinstance(self.__copy, bool):
+            raise TypeError(f"copy must be bool. Got {type(self.__copy)}")
+        if self.__copy is False:
+            raise NotImplementedError("copy=False not supported yet in Pandas API")
+
+        if self.__errors not in ("ignore", "raise"):
+            raise ValueError(f"errors must be 'ignore' or 'raise'. Got {self.__errors}")
+
+        # axis validation
+        if self.__axis not in (1, "columns", 0, "index"):
+            raise ValueError(f"Unsupported axis {self.__axis}")
+        if self.__axis in (0, "index"):
+            raise NotImplementedError("Renaming index not supported yet in Pandas API")
+
+        # index
+        if self.__index is not None:
+            raise NotImplementedError("Index mapper not supported yet in Pandas API")
+
+        # conflict validation
+        if self.__mapper and self.__columns:
+            raise ValueError("Cannot specify both 'axis' and any of 'index' or 'columns'")
+
+        self.__resolve_columns_mapping()  # runs validation
+        return True

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

@@ -1,120 +1,151 @@
-# 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 datetime import date
-from pylegend._typing import (
-    PyLegendList,
-    PyLegendSequence,
-    PyLegendUnion,
-)
-from pylegend.core.sql.metamodel import LongLiteral, QuerySpecification
-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 create_sub_query, copy_query
-from pylegend.core.tds.tds_column import TdsColumn
-from pylegend.core.tds.tds_frame import FrameToPureConfig, FrameToSqlConfig
-
-
-class TruncateFunction(PandasApiAppliedFunction):
-    __base_frame: PandasApiBaseTdsFrame
-    __before: int
-    __after: PyLegendUnion[int, None]
-    __axis: PyLegendUnion[str, int]
-    __copy: bool
-
-    @classmethod
-    def name(cls) -> str:
-        return "truncate"  # pragma: no cover
-
-    def __init__(
-            self,
-            base_frame: PandasApiBaseTdsFrame,
-            before: PyLegendUnion[date, str, int, None],
-            after: PyLegendUnion[date, str, int, None],
-            axis: PyLegendUnion[str, int],
-            copy: bool,
-    ) -> None:
-        self.__base_frame = base_frame
-        self.__before_input = before
-        self.__after_input = after
-        self.__axis = axis
-        self.__copy = copy
-
-    def to_sql(self, config: FrameToSqlConfig) -> QuerySpecification:
-        base_query: QuerySpecification = self.__base_frame.to_sql_query_object(config)
-        should_create_sub_query = (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)
-        new_query.offset = LongLiteral(self.__before)
-        if self.__after is not None:
-            new_query.limit = LongLiteral(self.__after - self.__before + 1)
-        return new_query
-
-    def to_pure(self, config: FrameToPureConfig) -> str:
-        if self.__after is None:
-            return f"{self.__base_frame.to_pure(config)}{config.separator(1)}" f"->drop({self.__before})"
-
-        start_row = self.__before
-        end_row = self.__after + 1
-        return f"{self.__base_frame.to_pure(config)}{config.separator(1)}" f"->slice({start_row}, {end_row})"
-
-    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:
-        if self.__axis not in [0, "index"]:
-            raise NotImplementedError(
-                f"The 'axis' parameter of the truncate function must be 0 or 'index', but got: {self.__axis}"
-            )
-
-        if self.__copy not in [True]:
-            raise NotImplementedError(f"The 'copy' parameter of the truncate function must be True, but got: {self.__copy}")
-
-        if self.__before_input is None:
-            self.__before = 0
-        else:
-            self.__before = self.get_positive_integer_or_raise_exception(self.__before_input, variable_name="before")
-
-        if self.__after_input is None:
-            self.__after = None
-            return True
-        else:
-            self.__after = self.get_positive_integer_or_raise_exception(self.__after_input, variable_name="after")
-
-            if self.__before > self.__after:
-                raise ValueError(
-                    f"The 'before' parameter of the truncate function must be less than or equal to the 'after' parameter, "
-                    f"but got: before={self.__before}, after={self.__after}"
-                )
-
-        return True
-
-    def get_positive_integer_or_raise_exception(
-        self, variable: PyLegendUnion[date, str, int, None], variable_name: str
-    ) -> int:
-        if type(variable) is not int:
-            raise NotImplementedError(
-                f"The '{variable_name}' parameter of the truncate function must be an integer, "
-                f"but got: {variable} (type: {type(variable).__name__})"
-            )
-
-        if variable < 0:
-            return 0
-        return variable
+# 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 datetime import date
+from pylegend._typing import (
+    PyLegendList,
+    PyLegendSequence,
+    PyLegendUnion,
+    PyLegendTuple,
+    PyLegendOptional,
+)
+from pylegend.core.sql.metamodel import LongLiteral, QuerySpecification
+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 create_sub_query, copy_query
+from pylegend.core.tds.tds_column import TdsColumn
+from pylegend.core.tds.tds_frame import FrameToPureConfig, FrameToSqlConfig
+
+
+class TruncateFunction(PandasApiAppliedFunction):
+    __base_frame: PandasApiBaseTdsFrame
+    __before: int
+    __after: PyLegendUnion[int, None]
+    __axis: PyLegendUnion[str, int]
+    __copy: bool
+
+    @classmethod
+    def name(cls) -> str:
+        return "truncate"  # pragma: no cover
+
+    def __init__(
+            self,
+            base_frame: PandasApiBaseTdsFrame,
+            before: PyLegendUnion[date, str, int, None],
+            after: PyLegendUnion[date, str, int, None],
+            axis: PyLegendUnion[str, int],
+            copy: bool,
+    ) -> None:
+        self.__base_frame = base_frame
+        self.__before_input = before
+        self.__after_input = after
+        self.__axis = axis
+        self.__copy = copy
+
+    def to_sql(self, config: FrameToSqlConfig) -> QuerySpecification:
+        base_query: QuerySpecification = self.__base_frame.to_sql_query_object(config)
+        should_create_sub_query = (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)
+        new_query.offset = LongLiteral(self.__before)
+
+        if self.__after is not None:
+            new_query.limit = LongLiteral(self.__after - self.__before + 1)
+        return new_query
+
+    def to_pure(self, config: FrameToPureConfig) -> str:
+        if self.__after is None:
+            return f"{self.__base_frame.to_pure(config)}{config.separator(1)}" f"->drop({self.__before})"
+
+        start_row = self.__before
+        end_row = self.__after + 1
+        return f"{self.__base_frame.to_pure(config)}{config.separator(1)}" f"->slice({start_row}, {end_row})"
+
+    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:
+        if self.__axis not in [0, "index"]:
+            raise NotImplementedError(
+                f"The 'axis' parameter of the truncate function must be 0 or 'index', but got: {self.__axis}"
+            )
+
+        if self.__copy not in [True]:
+            raise NotImplementedError(f"The 'copy' parameter of the truncate function must be True, but got: {self.__copy}")
+
+        self.__before, self.__after = self.__normalize_before_and_after(self.__before_input, self.__after_input)
+        return True
+
+    @staticmethod
+    def __normalize_before_and_after(
+            before_input: PyLegendUnion[date, str, int, None],
+            after_input: PyLegendUnion[date, str, int, None]
+    ) -> PyLegendTuple[int, PyLegendOptional[int]]:
+
+        if isinstance(before_input, (date, str)):
+            raise NotImplementedError(
+                f"The 'before' parameter of the truncate function must be of type integer or None, "
+                f"but got: before={before_input} (type: {type(before_input).__name__})")
+
+        if isinstance(after_input, (date, str)):
+            raise NotImplementedError(
+                f"The 'after' parameter of the truncate function must be of type integer or None, "
+                f"but got: after={after_input} (type: {type(after_input).__name__})")
+
+        def __raise_error_if_before_gt_after(before_input: int, after_input: int) -> None:
+            if before_input > after_input:
+                raise ValueError(
+                    f"The 'before' parameter of the truncate function must be less than or equal to the 'after' parameter, "
+                    f"but got: before={before_input}, after={after_input}")
+
+        if before_input is None:
+            if after_input is None:
+                return 0, None
+
+            if isinstance(after_input, int) and after_input >= 0:
+                return 0, after_input
+
+            if isinstance(after_input, int) and after_input < 0:
+                return 0, -1
+
+        if isinstance(before_input, int) and before_input >= 0:
+            if after_input is None:
+                return before_input, None
+
+            if isinstance(after_input, int) and after_input >= 0:
+                __raise_error_if_before_gt_after(before_input, after_input)
+                return before_input, after_input
+
+            if isinstance(after_input, int) and after_input < 0:
+                __raise_error_if_before_gt_after(before_input, after_input)
+
+        if isinstance(before_input, int) and before_input < 0:
+            if after_input is None:
+                return 0, None
+
+            if isinstance(after_input, int) and after_input >= 0:
+                __raise_error_if_before_gt_after(before_input, after_input)
+                return 0, after_input
+
+            if isinstance(after_input, int) and after_input < 0:
+                __raise_error_if_before_gt_after(before_input, after_input)
+                return 0, -1
+
+        return 0, 0  # pragma: no cover

pylegend/core/tds/pandas_api/frames/pandas_api_applied_function_tds_frame.py

@@ -13,16 +13,17 @@
 # limitations under the License.
 
 from abc import ABCMeta, abstractmethod
+
 from pylegend._typing import (
     PyLegendSequence,
     PyLegendList,
+    PyLegendType
 )
 from pylegend.core.sql.metamodel import QuerySpecification
+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 FrameToSqlConfig
 from pylegend.core.tds.tds_frame import FrameToPureConfig
-from pylegend.core.tds.pandas_api.frames.pandas_api_base_tds_frame import PandasApiBaseTdsFrame
-
+from pylegend.core.tds.tds_frame import FrameToSqlConfig, PyLegendTdsFrame
 
 __all__: PyLegendSequence[str] = [
     "PandasApiAppliedFunctionTdsFrame",
@@ -69,6 +70,9 @@ class PandasApiAppliedFunctionTdsFrame(PandasApiBaseTdsFrame):
         super().__init__(columns=applied_function.calculate_columns())
         self.__applied_function = applied_function
 
+    def get_super_type(self) -> PyLegendType[PyLegendTdsFrame]:
+        return type(self)  # pragma: no cover
+
     def to_sql_query_object(self, config: FrameToSqlConfig) -> QuerySpecification:
         return self.__applied_function.to_sql(config)