handle other iterables than list in queries (#257)

CHANGELOG.rst

@@ -1,6 +1,14 @@
 Changelog
 =========
 
+Version v3.0.1
+--------------
+
+Bug Fixes
+~~~~~~~~~
+- Fixed a bug causing some iterables (e.g., tuples) in queries not to work as expected
+
+
 Version v3.0.0
 --------------
 

bluepysnap/query.py

@@ -126,7 +126,7 @@ def _properties_mask(data, population_name, queries):
             prop_mask = np.logical_and(prop >= v1, prop <= v2)
         elif isinstance(values, Mapping):
             prop_mask = _complex_query(prop, values)
-        elif isinstance(values, list):
+        elif utils.is_iterable(values):
             prop_mask = prop.isin(values)
         else:
             prop_mask = prop == values

tests/test_edge_population.py

@@ -206,6 +206,11 @@ def test_ids(self):
         # if sample > population.size --> sample = population.size
         npt.assert_equal(len(self.test_obj.ids(sample=25)), 4)
 
+        # check iterables in queries
+        npt.assert_equal(self.test_obj.ids({"@target_node": [0, 1]}), [0, 1, 2, 3])
+        npt.assert_equal(self.test_obj.ids({"@target_node": (0, 1)}), [0, 1, 2, 3])
+        npt.assert_equal(self.test_obj.ids({"@target_node": map(int, [0, 1])}), [0, 1, 2, 3])
+
     def test_get_1(self):
         properties = [
             Synapse.PRE_GID,

tests/test_node_population.py

@@ -215,6 +215,8 @@ def test_ids(self):
         # same query with a $and operator
         npt.assert_equal(_call({"$and": [{Cell.MTYPE: "L6_Y"}, {Cell.MORPHOLOGY: "morph-B"}]}), [1])
         npt.assert_equal(_call({Cell.MORPHOLOGY: ["morph-A", "morph-B"]}), [0, 1])
+        npt.assert_equal(_call({Cell.MORPHOLOGY: ("morph-A", "morph-B")}), [0, 1])
+        npt.assert_equal(_call({Cell.MORPHOLOGY: map(str, ["morph-A", "morph-B"])}), [0, 1])
         npt.assert_equal(_call({"$and": [{}, {}]}), [0, 1, 2])
         npt.assert_equal(_call({"$and": [{}, {Cell.MORPHOLOGY: "morph-B"}]}), [1])
         # same query with a $or operator

tests/test_query.py

@@ -125,6 +125,18 @@ def test_resolve_ids():
     assert [False, False, False] == resolve_ids(data, "", {"$or": []}).tolist()
     assert [True, True, True] == resolve_ids(data, "", {"$and": []}).tolist()
 
+    # Test that iterables are handled correctly
+    expected = [False, True, True]
+    iterables = [
+        ["eight", "nine"],
+        ("eight", "nine"),
+        map(str, ["eight", "nine"]),
+        pd.Series(["eight", "nine"]),
+        {"eight", "nine"},
+    ]
+    for it in iterables:
+        assert expected == resolve_ids(data, "", {"str": it}).tolist()
+
     with pytest.raises(BluepySnapError) as e:
         resolve_ids(data, "", {"str": {"$regex": "*.some", "edge_id": 2}})
     assert "Value operators can't be used with plain values" in e.value.args[0]