tensor: Add coverage for improved indexing capability

Author: ntjohnson1

pyttb/tensor.py

@@ -1276,7 +1276,6 @@ def __setitem__(self, key, value):
         """
         # Figure out if we are doing a subtensor, a list of subscripts or a list of
         # linear indices
-        # print(f"Key: {key} {type(key)}")
         access_type = "error"
         # TODO pull out this big decision tree into a function
         if isinstance(key, (float, int, np.generic, slice)):
@@ -1296,15 +1295,11 @@ def __setitem__(self, key, value):
                 validSubtensor = [
                     isinstance(keyElement, (int, slice, Iterable)) for keyElement in key
                 ]
-                # TODO probably need to confirm the Iterable is in fact numeric
                 if np.all(validSubtensor):
                     access_type = "subtensor"
             elif isinstance(key, Iterable):
                 key = np.array(key)
-                # Clean up copy paste
-                if len(key.shape) > 1 and key.shape[1] >= self.ndims:
-                    access_type = "subscripts"
-                elif len(key.shape) == 1 or key.shape[1] == 1:
+                if len(key.shape) == 1 or key.shape[1] == 1:
                     access_type = "linear indices"
 
         # Case 1: Rectangular Subtensor
@@ -1351,6 +1346,12 @@ def _set_subtensor(self, key, value):
                 else:
                     sliceCheck.append(element.stop)
             elif isinstance(element, Iterable):
+                if any(
+                    not isinstance(entry, (float, int, np.generic)) for entry in element
+                ):
+                    raise ValueError(
+                        f"Entries for setitem must be numeric but recieved, {element}"
+                    )
                 sliceCheck.append(max(element))
             else:
                 sliceCheck.append(element)

tests/test_tensor.py

@@ -251,6 +251,13 @@ def test_tensor__setitem__(sample_tensor_2way):
     # Subtensor add dimension
     empty_tensor[0, 0, 0] = 2
 
+    # Subtensor with lists
+    some_tensor = ttb.tenones((3, 3))
+    some_tensor[[0, 1], [0, 1]] = 11
+    assert some_tensor[0, 0] == 11
+    assert some_tensor[1, 1] == 11
+    assert np.all(some_tensor[[0, 1], [0, 1]].data == 11)
+
     # Subscripts with constant
     tensorInstance[np.array([[1, 1]])] = 13.0
     dataGrowth[1, 1] = 13.0
@@ -293,6 +300,13 @@ def test_tensor__setitem__(sample_tensor_2way):
     dataGrowth[np.unravel_index([0, 3, 4], dataGrowth.shape, "F")] = 13
     assert (tensorInstance.data == dataGrowth).all()
 
+    # Linear index with multiple indicies
+    some_tensor = ttb.tenones((3, 3))
+    some_tensor[[0, 1]] = 2
+    assert some_tensor[0] == 2
+    assert some_tensor[1] == 2
+    assert np.array_equal(some_tensor[[0, 1]], [2, 2])
+
     # Test Empty Tensor Set Item, subtensor
     emptyTensor = ttb.tensor.from_data(np.array([]))
     emptyTensor[0, 0, 0] = 0
@@ -313,10 +327,17 @@ def test_tensor__setitem__(sample_tensor_2way):
     )
 
     # Attempting to set some other way
-    # TODO either catch this error ourselves or specify more specific exception we expect here
-    with pytest.raises(Exception) as excinfo:
+    with pytest.raises(ValueError) as excinfo:
         tensorInstance[0, "a", 5] = 13.0
-    # assert "Invalid use of tensor setitem" in str(excinfo)
+    assert "must be numeric" in str(excinfo)
+
+    with pytest.raises(AssertionError) as excinfo:
+
+        class BadKey:
+            pass
+
+        tensorInstance[BadKey] = 13.0
+    assert "Invalid use of tensor setitem" in str(excinfo)
 
 
 @pytest.mark.indevelopment
@@ -346,6 +367,11 @@ def test_tensor__getitem__(sample_tensor_2way):
         tensorInstance[np.array([[0, 0], [1, 1]]), "extract"]
         == params["data"][([0, 0], [1, 1])]
     ).all()
+    # Case 2a: Extract doesn't seem to be needed
+    assert tensorInstance[np.array([0, 0])] == params["data"][0, 0]
+    assert (
+        tensorInstance[np.array([[0, 0], [1, 1]])] == params["data"][([0, 0], [1, 1])]
+    ).all()
 
     # Case 2b: Linear Indexing
     assert tensorInstance[np.array([0])] == params["data"][0, 0]