diff options
Diffstat (limited to 'polymatrix/expression/init/initgetitemexpr.py')
-rw-r--r-- | polymatrix/expression/init/initgetitemexpr.py | 14 |
1 files changed, 12 insertions, 2 deletions
diff --git a/polymatrix/expression/init/initgetitemexpr.py b/polymatrix/expression/init/initgetitemexpr.py index 140fa3a..5fea7a5 100644 --- a/polymatrix/expression/init/initgetitemexpr.py +++ b/polymatrix/expression/init/initgetitemexpr.py @@ -1,12 +1,22 @@ +from numpy import isin from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin from polymatrix.expression.impl.getitemexprimpl import GetItemExprImpl def init_get_item_expr( underlying: ExpressionBaseMixin, - index: tuple, + index: tuple[tuple[int, ...], tuple[int, ...]], ): + + def get_hashable_slice(index): + if isinstance(index, slice): + return GetItemExprImpl.Slice(start=index.start, stop=index.stop, step=index.step) + else: + return index + + proper_index = (get_hashable_slice(index[0]), get_hashable_slice(index[1])) + return GetItemExprImpl( underlying=underlying, - index=index, + index=proper_index, ) |