From 03167f807de0a3a932b9a8e6d374a99970265673 Mon Sep 17 00:00:00 2001 From: Giulio Romualdi Date: Thu, 9 Jan 2025 11:38:01 +0100 Subject: [PATCH] Enhance __setitem__ and __getitem__ methods to support Ellipsis indexing in CasadiLike class --- src/adam/casadi/casadi_like.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/src/adam/casadi/casadi_like.py b/src/adam/casadi/casadi_like.py index 9feac02..ed01213 100644 --- a/src/adam/casadi/casadi_like.py +++ b/src/adam/casadi/casadi_like.py @@ -85,11 +85,27 @@ def __truediv__(self, other: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike" def __setitem__(self, idx, value: Union["CasadiLike", npt.ArrayLike]): """Overrides set item operator""" - self.array[idx] = value.array if type(self) is type(value) else value + if idx is Ellipsis: + self.array = value.array if isinstance(value, CasadiLike) else value + elif isinstance(idx, tuple) and Ellipsis in idx: + idx = tuple(slice(None) if i is Ellipsis else i for i in idx) + self.array[idx] = value.array if isinstance(value, CasadiLike) else value + else: + self.array[idx] = value.array if isinstance(value, CasadiLike) else value + def __getitem__(self, idx) -> "CasadiLike": """Overrides get item operator""" - return CasadiLike(self.array[idx]) + if idx is Ellipsis: + # Handle the case where only Ellipsis is passed + return CasadiLike(self.array) + elif isinstance(idx, tuple) and Ellipsis in idx: + # Handle the case where Ellipsis is part of a tuple + idx = tuple(slice(None) if k is Ellipsis else k for k in idx) + return CasadiLike(self.array[idx]) + else: + # For other cases, delegate to the CasADi object's __getitem__ + return CasadiLike(self.array[idx]) @property def T(self) -> "CasadiLike":