Merge pull request #362 from psobot/psobot/fix-crash-on-get-items

Throw exception instead of segfaulting when passing a scalar to get_items.

Author: yurymalkov

python_bindings/bindings.cpp

@@ -260,11 +260,16 @@ class Index {
         if (!ids_.is_none()) {
             py::array_t < size_t, py::array::c_style | py::array::forcecast > items(ids_);
             auto ids_numpy = items.request();
-            std::vector<size_t> ids1(ids_numpy.shape[0]);
-            for (size_t i = 0; i < ids1.size(); i++) {
-                ids1[i] = items.data()[i];
+
+            if (ids_numpy.ndim == 0) {
+              throw std::invalid_argument("get_items accepts a list of indices and returns a list of vectors");
+            } else {
+              std::vector<size_t> ids1(ids_numpy.shape[0]);
+              for (size_t i = 0; i < ids1.size(); i++) {
+                  ids1[i] = items.data()[i];
+              }
+              ids.swap(ids1);
             }
-            ids.swap(ids1);
         }
 
         std::vector<std::vector<data_t>> data;

python_bindings/tests/bindings_test_getdata.py

@@ -41,6 +41,9 @@ def testGettingItems(self):
         print("Adding all elements (%d)" % (len(data)))
         p.add_items(data, labels)
 
+        # Getting data by label should raise an exception if a scalar is passed:
+        self.assertRaises(ValueError, lambda: p.get_items(labels[0]))
+
         # After adding them, all labels should be retrievable
         returned_items = p.get_items(labels)
         self.assertSequenceEqual(data.tolist(), returned_items)