NonUniformBatch: implement __getitem__ to allow index array (batch[idx]).

Author: PMeira

.gitignore

@@ -131,3 +131,6 @@ altdss/_version.py
 
 apidocs
 electricdss-tst
+htmlcov/
+.coverage
+private_tests/

altdss/Batch.py

@@ -648,5 +648,24 @@ def __iter__(self) -> Iterator[DSSObj]:
         for idx in range(cnt[0]):
             yield self._obj_cls(self._api_util, ptr[idx])
 
+    def __getitem__(self, idx: int) -> DSSObj:
+        '''
+        Return an object of the batch by index (0-based).
+        '''
+        _pointer, _count = self._get_ptr_cnt()
+        if idx > _count or idx < 0:
+            raise IndexError('Invalid object index inside the batch')
+        
+        ptr = _pointer[idx]
+        if ptr == self._ffi.NULL:
+            return None
+
+        if self._obj_cls is None:
+            cls_idx = self._lib.Obj_GetClassIdx(ptr)
+            pycls = DSSObj._idx_to_cls[cls_idx]
+            return pycls(self._api_util, ptr)
+
+        return self._obj_cls(self._api_util, ptr)
+
 
 __all__ = ('DSSBatch', 'NonUniformBatch', )

docs/changelog.md

@@ -6,6 +6,7 @@ relevant. See [DSS C-API's repository](https://github.com/dss-extensions/dss_cap
 ## 0.2.2
 
 - CircuitElementBatch: fix `MaxCurrent`. This will require the backend to be updated to v0.14.3.
+- NonUniformBatch: allow `batch[idx]` to get a single element by index.
 
 ## 0.2.1