Sparse2Corpus: update __getitem__ to work on slices, lists and ellipsis

gensim/matutils.py

@@ -597,23 +597,24 @@ def __iter__(self):
     def __len__(self):
         return self.sparse.shape[1]
 
-    def __getitem__(self, document_index):
+    def __getitem__(
+        self, key: Union[int, List[int], type(...), slice]
+    ) -> Sparse2Corpus:
         """Retrieve a document vector from the corpus by its index.
 
         Parameters
         ----------
-        document_index : int
-            Index of document
+        key
+            Index of document or slice for documents
 
         Returns
         -------
-        list of (int, number)
-            Document in BoW format.
-
+        Sparse2Corpus with data subset
         """
-        indprev = self.sparse.indptr[document_index]
-        indnow = self.sparse.indptr[document_index + 1]
-        return list(zip(self.sparse.indices[indprev:indnow], self.sparse.data[indprev:indnow]))
+        if not isinstance(key, (int, list, type(...), slice)):
+            raise TypeError(f"Indexing by type {type(key)} not supported.")
+        sparse = self.sparse.__getitem__((slice(None, None, None), key))
+        return Sparse2Corpus(sparse)
 
 
 def veclen(vec):