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

pylegend/core/sql/metamodel_extension.py

@@ -20,6 +20,7 @@ from pylegend._typing import (
 from pylegend.core.sql.metamodel import (
     Expression,
     Window,
+    StringLiteral,
 )
 
 __all__: PyLegendSequence[str] = [
@@ -81,7 +82,12 @@ __all__: PyLegendSequence[str] = [
     "EpochExpression",
     "WindowExpression",
     "ConstantExpression",
-    "StringSubStringExpression"
+    "StringSubStringExpression",
+    "DateAdjustExpression",
+    "BitwiseNotExpression",
+    "DateDiffExpression",
+    "DateTimeBucketExpression",
+    "DateType"
 ]
 
 
@@ -763,3 +769,73 @@ class StringSubStringExpression(Expression):
         self.value = value
         self.start = start
         self.end = end
+
+
+class DateAdjustExpression(Expression):
+    date: "Expression"
+    number: "Expression"
+    duration_unit: "StringLiteral"
+
+    def __init__(
+            self,
+            date: "Expression",
+            number: "Expression",
+            duration_unit: "StringLiteral",
+    ) -> None:
+        super().__init__(_type="dateAdjustExpression")
+        self.date = date
+        self.number = number
+        self.duration_unit = duration_unit
+
+
+class DateDiffExpression(Expression):
+    start_date: "Expression"
+    end_date: "Expression"
+    duration_unit: "StringLiteral"
+
+    def __init__(
+            self,
+            start_date: "Expression",
+            end_date: "Expression",
+            duration_unit: "StringLiteral",
+    ) -> None:
+        super().__init__(_type="dateDiffExpression")
+        self.start_date = start_date
+        self.end_date = end_date
+        self.duration_unit = duration_unit
+
+
+class DateType(Enum):
+    DateTime = 1
+    StrictDate = 2
+
+
+class DateTimeBucketExpression(Expression):
+    date: "Expression"
+    quantity: "Expression"
+    duration_unit: "StringLiteral"
+    date_type: DateType
+
+    def __init__(
+            self,
+            date: "Expression",
+            quantity: "Expression",
+            duration_unit: "StringLiteral",
+            date_type: DateType = DateType.DateTime,
+    ) -> None:
+        super().__init__(_type="dateTimeBucketExpression")
+        self.date = date
+        self.quantity = quantity
+        self.duration_unit = duration_unit
+        self.date_type = date_type
+
+
+class BitwiseNotExpression(Expression):
+    value: "Expression"
+
+    def __init__(
+            self,
+            value: "Expression",
+    ) -> None:
+        super().__init__(_type="bitwiseNotExpression")
+        self.value = value

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

@@ -100,7 +100,7 @@ class AggregateFunction(PandasApiAppliedFunction):
 
         if isinstance(self.__base_frame, PandasApiGroupbyTdsFrame):
             columns_to_retain: PyLegendList[str] = [
-                db_extension.quote_identifier(x) for x in self.__base_frame.grouping_column_name_list()
+                db_extension.quote_identifier(x.get_name()) for x in self.__base_frame.get_grouping_columns()
             ]
             new_cols_with_index: PyLegendList[PyLegendTuple[int, "SelectItem"]] = []
             for col in new_query.select.selectItems:
@@ -127,8 +127,8 @@ class AggregateFunction(PandasApiAppliedFunction):
         if isinstance(self.__base_frame, PandasApiGroupbyTdsFrame):
             tds_row = PandasApiTdsRow.from_tds_frame("r", self.base_frame())
             new_query.groupBy = [
-                (lambda x: x[c])(tds_row).to_sql_expression({"r": new_query}, config)
-                for c in self.__base_frame.grouping_column_name_list()
+                (lambda x: x[c.get_name()])(tds_row).to_sql_expression({"r": new_query}, config)
+                for c in self.__base_frame.get_grouping_columns()
             ]
 
         return new_query
@@ -149,8 +149,8 @@ class AggregateFunction(PandasApiAppliedFunction):
 
         if isinstance(self.__base_frame, PandasApiGroupbyTdsFrame):
             group_strings = []
-            for col_name in self.__base_frame.grouping_column_name_list():
-                group_strings.append(escape_column_name(col_name))
+            for col in self.__base_frame.get_grouping_columns():
+                group_strings.append(escape_column_name(col.get_name()))
 
             pure_expression = (
                 f"{self.base_frame().to_pure(config)}{config.separator(1)}" + f"->groupBy({config.separator(2)}"
@@ -182,7 +182,8 @@ class AggregateFunction(PandasApiAppliedFunction):
 
         if isinstance(self.__base_frame, PandasApiGroupbyTdsFrame):
             base_cols_map = {c.get_name(): c for c in self.base_frame().columns()}
-            for group_col_name in self.__base_frame.grouping_column_name_list():
+            for group_col in self.__base_frame.get_grouping_columns():
+                group_col_name = group_col.get_name()
                 if group_col_name in base_cols_map:
                     new_columns.append(base_cols_map[group_col_name].copy())
 
@@ -231,6 +232,10 @@ class AggregateFunction(PandasApiAppliedFunction):
 
         tds_row = PandasApiTdsRow.from_tds_frame("r", self.base_frame())
 
+        group_cols: set[str] = set()
+        if isinstance(self.__base_frame, PandasApiGroupbyTdsFrame):
+            group_cols = set([col.get_name() for col in self.__base_frame.get_grouping_columns()])
+
         for column_name, agg_input in normalized_func.items():
             mapper_function: PyLegendCallable[[PandasApiTdsRow], PyLegendPrimitiveOrPythonPrimitive] = eval(
                 f'lambda r: r["{column_name}"]'
@@ -259,7 +264,12 @@ class AggregateFunction(PandasApiAppliedFunction):
                 normalized_agg_func = self.__normalize_agg_func_to_lambda_function(agg_input)
                 agg_result = normalized_agg_func(collection)
 
-                self.__aggregates_list.append((column_name, map_result, agg_result))
+                if column_name in group_cols:
+                    alias = self._generate_column_alias(column_name, agg_input, 0)
+                else:
+                    alias = column_name
+
+                self.__aggregates_list.append((alias, map_result, agg_result))
 
         return True
 
@@ -274,13 +284,13 @@ class AggregateFunction(PandasApiAppliedFunction):
         all_cols = [col.get_name() for col in self.base_frame().columns()]
 
         if isinstance(self.__base_frame, PandasApiGroupbyTdsFrame):
-            group_cols = set(self.__base_frame.grouping_column_name_list())
+            group_cols = set([col.get_name() for col in self.__base_frame.get_grouping_columns()])
 
-            selected_cols = self.__base_frame.selected_columns()
+            selected_cols = self.__base_frame.get_selected_columns()
 
             if selected_cols is not None:
-                validation_columns = selected_cols
-                default_broadcast_columns = selected_cols
+                validation_columns = [col.get_name() for col in selected_cols]
+                default_broadcast_columns = [col.get_name() for col in selected_cols]
             else:
                 validation_columns = all_cols
                 default_broadcast_columns = [c for c in all_cols if c not in group_cols]

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

@@ -0,0 +1,99 @@
+# 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
+from pylegend._typing import (
+    PyLegendUnion,
+    PyLegendTuple,
+    PyLegendSequence,
+)
+
+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] = [
+    "PandasApiIlocIndexer"
+]
+
+
+class PandasApiIlocIndexer:
+    _frame: "PandasApiBaseTdsFrame"
+
+    def __init__(self, frame: "PandasApiBaseTdsFrame") -> None:
+        self._frame = frame
+
+    def __getitem__(  # type: ignore
+            self,
+            key: PyLegendUnion[int, slice, PyLegendTuple[PyLegendUnion[int, slice], ...]]
+    ) -> "PandasApiTdsFrame":
+        if isinstance(key, tuple):
+            if len(key) > 2:
+                raise IndexError("Too many indexers")
+            elif len(key) == 1:
+                rows, cols = key[0], slice(None, None, None)
+            else:
+                rows, cols = key  # type: ignore
+        else:
+            rows, cols = key, slice(None, None, None)
+
+        # Row selection
+        row_frame = self._handle_row_selection(rows)
+
+        # Column selection
+        return self._handle_column_selection(row_frame, cols)
+
+    def _handle_row_selection(self, rows: PyLegendUnion[int, slice]) -> "PandasApiTdsFrame":  # type: ignore
+        if isinstance(rows, slice):
+            if rows.step is not None and rows.step != 1:
+                raise NotImplementedError("iloc with slice step other than 1 is not supported yet in Pandas Api")
+
+            start = rows.start
+            stop = rows.stop
+            after = stop - 1 if stop is not None else None
+            return self._frame.truncate(before=start, after=after)
+
+        elif isinstance(rows, int):
+            return self._frame.truncate(before=rows, after=rows)
+
+        else:
+            raise NotImplementedError(
+                f"iloc supports integer, slice, or tuple of these, but got indexer of type: {type(rows)}"
+            )
+
+    def _handle_column_selection(  # type: ignore
+            self,
+            frame: "PandasApiTdsFrame",
+            cols: PyLegendUnion[int, slice]
+    ) -> "PandasApiTdsFrame":
+        if isinstance(cols, slice):
+            if cols.step is not None and cols.step != 1:
+                raise NotImplementedError("iloc with slice step other than 1 is not supported yet in Pandas Api")
+
+            all_columns = [c.get_name() for c in frame.columns()]
+            selected_columns = all_columns[cols]
+            return frame.filter(items=selected_columns)
+
+        elif isinstance(cols, int):
+            all_columns = [c.get_name() for c in frame.columns()]
+            if not -len(all_columns) <= cols < len(all_columns):
+                raise IndexError("single positional indexer is out-of-bounds")
+            selected_column = all_columns[cols]
+            return frame.filter(items=[selected_column])
+
+        else:
+            raise NotImplementedError(
+                f"iloc supports integer, slice, or tuple of these, but got indexer of type: {type(cols)}"
+            )

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/pandas_api_applied_function_tds_frame.py

@@ -85,3 +85,6 @@ class PandasApiAppliedFunctionTdsFrame(PandasApiBaseTdsFrame):
             for x in [self.__applied_function.base_frame()] + self.__applied_function.tds_frame_parameters()
             for y in x.get_all_tds_frames()
         ] + [self]
+
+    def get_applied_function(self) -> PandasApiAppliedFunction:
+        return self.__applied_function

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

@@ -15,7 +15,7 @@
 import copy
 from abc import ABCMeta, abstractmethod
 from datetime import date, datetime
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
 from typing_extensions import Concatenate
 
@@ -69,6 +69,8 @@ from pylegend.extensions.tds.result_handler import (
 if TYPE_CHECKING:
     from pylegend.core.language.pandas_api.pandas_api_series import Series
     from pylegend.core.tds.pandas_api.frames.pandas_api_groupby_tds_frame import PandasApiGroupbyTdsFrame
+    from pylegend.core.tds.pandas_api.frames.functions.iloc import PandasApiIlocIndexer
+    from pylegend.core.tds.pandas_api.frames.functions.loc import PandasApiLocIndexer
 
 __all__: PyLegendSequence[str] = [
     "PandasApiBaseTdsFrame"
@@ -94,6 +96,14 @@ class PandasApiBaseTdsFrame(PandasApiTdsFrame, BaseTdsFrame, metaclass=ABCMeta):
             return [c.copy() for c in self.__columns]
         return self._transformed_frame.columns()
 
+    @overload  # type: ignore[override]
+    def __getitem__(self, key: str) -> "Series":
+        ...
+
+    @overload
+    def __getitem__(self, key: PyLegendList[str]) -> "PandasApiTdsFrame":
+        ...
+
     def __getitem__(
             self,
             key: PyLegendUnion[str, PyLegendList[str], PyLegendBoolean]
@@ -663,7 +673,7 @@ class PandasApiBaseTdsFrame(PandasApiTdsFrame, BaseTdsFrame, metaclass=ABCMeta):
             # Compute row callable via func on the Series
             def _row_callable(
                     _row: PandasApiTdsRow,
-                    _s: Series = series,  # type: ignore
+                    _s: Series = series,
                     _a: PyLegendTuple[PyLegendPrimitiveOrPythonPrimitive, ...] = args,
                     _k: PyLegendPrimitiveOrPythonPrimitive = kwargs  # type: ignore
             ) -> PyLegendPrimitiveOrPythonPrimitive:
@@ -675,6 +685,44 @@ class PandasApiBaseTdsFrame(PandasApiTdsFrame, BaseTdsFrame, metaclass=ABCMeta):
             AssignFunction(self, col_definitions=col_definitions)  # type: ignore
         )
 
+    @property
+    def iloc(self) -> "PandasApiIlocIndexer":
+        """
+        Purely integer-location based indexing for selection by position.
+        .iloc[] is primarily integer position based (from 0 to length-1 of the axis).
+
+        Allowed inputs are:
+        - An integer, e.g. 5.
+        - A slice object with ints, e.g. 1:7.
+        - A tuple of row and column indexes, e.g., (slice(1, 5), slice(0, 2))
+
+        Other pandas iloc features such as list of integers, boolean arrays, and callables
+        are not supported and will raise a NotImplementedError.
+        """
+        from pylegend.core.tds.pandas_api.frames.functions.iloc import PandasApiIlocIndexer
+        return PandasApiIlocIndexer(self)
+
+    @property
+    def loc(self) -> "PandasApiLocIndexer":
+        """
+        Access a group of rows and columns by label(s) or a boolean array.
+        .loc[] is primarily label based, but may also be used with a boolean array.
+
+        Allowed inputs are:
+        - A single label, e.g. 5 or 'a', (note that 5 is interpreted as a
+          label of the index, not as an integer position along the index).
+        - A list or array of labels, e.g. ['a', 'b', 'c'].
+        - A slice object with labels, e.g. 'a':'f'.
+        - A boolean array of the same length as the axis being sliced.
+        - A callable function with one argument (the calling Series or
+          DataFrame) and that returns valid output for indexing (one of the above).
+
+        Currently, for row selection, only callable function or complete slice are supported.
+        For column selection, string labels, lists of string labels, and slices of string labels are supported.
+        """
+        from pylegend.core.tds.pandas_api.frames.functions.loc import PandasApiLocIndexer
+        return PandasApiLocIndexer(self)
+
     def head(self, n: int = 5) -> "PandasApiTdsFrame":
         """
         Return the first `n` rows by calling truncate on rows.