pylegend 0.11.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

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

@@ -0,0 +1,136 @@
+# Copyright 2026 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 typing import TYPE_CHECKING
+
+import pandas as pd
+
+from pylegend._typing import (
+    PyLegendUnion,
+    PyLegendSequence,
+    PyLegendTuple,
+    PyLegendCallable
+)
+from pylegend.core.language import PyLegendBoolean
+from pylegend.core.tds.pandas_api.frames.functions.filtering import PandasApiFilteringFunction
+from pylegend.core.tds.pandas_api.frames.pandas_api_applied_function_tds_frame import PandasApiAppliedFunctionTdsFrame
+
+if TYPE_CHECKING:
+    from pylegend.core.tds.pandas_api.frames.pandas_api_base_tds_frame import PandasApiBaseTdsFrame
+    from pylegend.core.tds.pandas_api.frames.pandas_api_tds_frame import PandasApiTdsFrame
+
+__all__: PyLegendSequence[str] = [
+    "PandasApiLocIndexer"
+]
+
+
+class PandasApiLocIndexer:
+    _frame: "PandasApiBaseTdsFrame"
+
+    def __init__(self, frame: "PandasApiBaseTdsFrame") -> None:
+        self._frame = frame
+
+    def __getitem__(  # type: ignore
+            self,
+            key: PyLegendUnion[
+                slice,
+                PyLegendBoolean,
+                PyLegendCallable[["PandasApiBaseTdsFrame"], PyLegendBoolean],
+                PyLegendTuple[
+                    PyLegendUnion[slice, PyLegendBoolean, PyLegendCallable[["PandasApiBaseTdsFrame"], PyLegendBoolean]],
+                    PyLegendUnion[str, slice, PyLegendSequence[str], PyLegendSequence[bool]]
+                ]
+            ]
+    ) -> "PandasApiTdsFrame":
+        rows: PyLegendUnion[  # type: ignore
+            slice,
+            PyLegendBoolean,
+            PyLegendCallable[["PandasApiBaseTdsFrame"], PyLegendBoolean]
+        ]
+        cols: PyLegendUnion[str, slice, PyLegendSequence[str], PyLegendSequence[bool]]  # type: ignore
+
+        if isinstance(key, tuple):
+            if len(key) == 1:
+                rows, cols = key[0], slice(None, None, None)
+            elif len(key) == 2:
+                rows, cols = key[0], key[1]
+            else:
+                raise IndexError("Too many indexers")
+        else:
+            rows, cols = key, slice(None, None, None)
+
+        row_frame = self._handle_row_selection(rows)
+        return self._handle_column_selection(row_frame, cols)
+
+    def _handle_row_selection(  # type: ignore
+            self,
+            rows: PyLegendUnion[slice, PyLegendBoolean, PyLegendCallable[["PandasApiBaseTdsFrame"], PyLegendBoolean]]
+    ) -> "PandasApiTdsFrame":
+        if isinstance(rows, slice):
+            if rows.start is None and rows.stop is None and rows.step is None:
+                return self._frame
+            else:
+                raise TypeError(
+                    "loc supports only ':' for row slicing. "
+                    "Label-based slicing for rows is not supported."
+                )
+
+        if isinstance(rows, PyLegendBoolean):
+            return PandasApiAppliedFunctionTdsFrame(
+                PandasApiFilteringFunction(self._frame, filter_expr=rows)
+            )
+
+        if callable(rows):
+            new_key = rows(self._frame)
+            return self._handle_row_selection(new_key)
+
+        raise TypeError(f"Unsupported key type for .loc row selection: {type(rows)}")
+
+    def _handle_column_selection(  # type: ignore
+            self,
+            frame: "PandasApiTdsFrame",
+            cols: PyLegendUnion[str, slice, PyLegendSequence[str], PyLegendSequence[bool]]
+    ) -> "PandasApiTdsFrame":
+        if isinstance(cols, slice) and cols.start is None and cols.stop is None and cols.step is None:
+            return frame
+
+        if isinstance(cols, str):
+            return frame.filter(items=[cols])
+
+        if isinstance(cols, (list, tuple)):
+            all_columns = [c.get_name() for c in frame.columns()]
+            is_boolean_list = all(isinstance(k, bool) for k in cols)
+
+            if is_boolean_list:
+                if len(cols) != len(all_columns):
+                    raise IndexError(f"Boolean index has wrong length: {len(cols)} instead of {len(all_columns)}")
+                selected_columns = [col for col, select in zip(all_columns, cols) if select]
+                return frame.filter(items=selected_columns)
+            else:
+                missing_cols = [c for c in cols if c not in all_columns]
+                if missing_cols:
+                    raise KeyError(f"{missing_cols} not in index")
+                return frame.filter(items=cols)  # type: ignore
+
+        if isinstance(cols, slice):
+            all_columns = [c.get_name() for c in frame.columns()]
+            pd_index = pd.Index(all_columns)
+
+            slicer = pd_index.slice_indexer(start=cols.start, end=cols.stop, step=cols.step)
+            selected_columns = pd_index[slicer].tolist()
+            if not selected_columns:
+                return frame.head(0)
+            return frame.filter(items=selected_columns)
+
+        raise TypeError(f"Unsupported key type for .loc column selection: {type(cols)}")

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)