fix #377 add setitem method to ETS class

Author: jhavl

roboticstoolbox/robot/ETS.py

@@ -37,7 +37,7 @@
 )
 from copy import deepcopy
 from roboticstoolbox.robot.ET import ET, ET2
-from typing import Union, overload, List, Set, Tuple
+from typing import Union, overload, List, Set, Tuple, TypeVar
 from typing_extensions import Literal as L
 from sys import version_info
 from roboticstoolbox.tools.types import ArrayLike, NDArray
@@ -51,6 +51,8 @@
 else:  # pragma: nocover
     c_property = property
 
+T = TypeVar("T", bound="BaseETS")
+
 
 class BaseETS(UserList):
     def __init__(self, *args):
@@ -435,38 +437,32 @@ def m(self) -> int:
         return self._m
 
     @overload
-    def data(self: "ETS") -> List[ET]:
-        ...  # pragma: nocover
+    def data(self: "ETS") -> List[ET]: ...  # pragma: nocover
 
     @overload
-    def data(self: "ETS2") -> List[ET2]:
-        ...  # pragma: nocover
+    def data(self: "ETS2") -> List[ET2]: ...  # pragma: nocover
 
     @property
     def data(self):
         return self._data
 
     @data.setter
     @overload
-    def data(self: "ETS", new_data: List[ET]):
-        ...  # pragma: nocover
+    def data(self: "ETS", new_data: List[ET]): ...  # pragma: nocover
 
     @data.setter
     @overload
-    def data(self: "ETS", new_data: List[ET2]):
-        ...  # pragma: nocover
+    def data(self: "ETS", new_data: List[ET2]): ...  # pragma: nocover
 
     @data.setter
     def data(self, new_data):
         self._data = new_data
 
     @overload
-    def pop(self: "ETS", i: int = -1) -> ET:
-        ...  # pragma: nocover
+    def pop(self: "ETS", i: int = -1) -> ET: ...  # pragma: nocover
 
     @overload
-    def pop(self: "ETS2", i: int = -1) -> ET2:
-        ...  # pragma: nocover
+    def pop(self: "ETS2", i: int = -1) -> ET2: ...  # pragma: nocover
 
     def pop(self, i=-1):
         """
@@ -506,12 +502,10 @@ def pop(self, i=-1):
         return item
 
     @overload
-    def split(self: "ETS") -> List["ETS"]:
-        ...  # pragma: nocover
+    def split(self: "ETS") -> List["ETS"]: ...  # pragma: nocover
 
     @overload
-    def split(self: "ETS2") -> List["ETS2"]:
-        ...  # pragma: nocover
+    def split(self: "ETS2") -> List["ETS2"]: ...  # pragma: nocover
 
     def split(self):
         """
@@ -547,15 +541,7 @@ def split(self):
 
         return segments
 
-    @overload
-    def inv(self: "ETS") -> "ETS":
-        ...  # pragma: nocover
-
-    @overload
-    def inv(self: "ETS2") -> "ETS2":
-        ...  # pragma: nocover
-
-    def inv(self):
+    def inv(self: T) -> T:
         r"""
         Inverse of ETS
 
@@ -589,20 +575,16 @@ def inv(self):
         return self.__class__([et.inv() for et in reversed(self.data)])
 
     @overload
-    def __getitem__(self: "ETS", i: int) -> ET:
-        ...  # pragma: nocover
+    def __getitem__(self: "ETS", i: int) -> ET: ...
 
     @overload
-    def __getitem__(self: "ETS", i: slice) -> List[ET]:
-        ...  # pragma: nocover
+    def __getitem__(self: "ETS", i: slice) -> List[ET]: ...
 
     @overload
-    def __getitem__(self: "ETS2", i: int) -> ET2:
-        ...  # pragma: nocover
+    def __getitem__(self: "ETS2", i: int) -> ET2: ...
 
     @overload
-    def __getitem__(self: "ETS2", i: slice) -> List[ET2]:
-        ...  # pragma: nocover
+    def __getitem__(self: "ETS2", i: slice) -> List[ET2]: ...
 
     def __getitem__(self, i):
         """
@@ -630,6 +612,41 @@ def __getitem__(self, i):
         """
         return self.data[i]  # can be [2] or slice, eg. [3:5]
 
+    @overload
+    def __setitem__(self: "ETS", i: int, value: ET): ...
+
+    @overload
+    def __setitem__(self: "ETS", i: slice, value: List[ET]): ...
+
+    @overload
+    def __setitem__(self: "ETS2", i: int, value: ET2): ...
+
+    @overload
+    def __setitem__(self: "ETS2", i: slice, value: List[ET2]): ...
+
+    def __setitem__(self, i, value):
+        """
+        Set an item in the ETS
+
+        Parameters
+        ----------
+        i
+            the index
+        value
+            the value to set
+
+        Examples
+        --------
+        .. runblock:: pycon
+        >>> from roboticstoolbox import ET
+        >>> e = ET.Rz() * ET.tx(1) * ET.Rz() * ET.tx(1)
+        >>> e[1] = ET.tx(2)
+        >>> e
+
+        """
+        self.data[i] = value
+        self._update_internals()
+
     def __deepcopy__(self, memo):
         new_data = []