enable selection init by dict

Author: jreadey

src/h5json/selections.py

@@ -65,8 +65,6 @@ def select(obj, args, fields=None):
     Indices, slices, ellipses, lists or boolean index arrays
         Returns a SimpleSelection instance with H5S_SEL_FANCY.
     """
-    if not isinstance(args, tuple):
-        args = (args,)
 
     if hasattr(obj, "shape"):
         obj_shape = obj.shape
@@ -75,6 +73,11 @@ def select(obj, args, fields=None):
     else:
         raise TypeError("Object must be a dataset or a shape tuple")
 
+    if isinstance(args, dict):
+        args = _handle_dict_selection(obj_shape, args)
+    elif not isinstance(args, tuple):
+        args = (args,)
+
     if len(obj_shape) == 0:
         # scalar object
         sel = ScalarSelection(obj_shape, args)
@@ -829,6 +832,57 @@ def translate(s1, s2):
     return select(s1.shape, tuple(args))
 
 
+def _handle_dict_selection(shape, arg):
+    """ Handle a dictionary-based selection, where the keys are dimension indices and
+    the values are slices or lists of coordinates. Returns a tuple of slices/lists for
+    each dimension, filling in full slices for unspecified dimensions.
+    """
+    rank = len(shape)
+    slices = []
+    if "start" not in arg:
+        start = (0,) * rank
+    else:
+        start = arg["start"]
+    if "stop" not in arg:
+        stop = shape
+    else:
+        stop = arg["stop"]
+    if "step" not in arg:
+        step = (1,) * rank
+    else:
+        step = arg["step"]
+
+    if isinstance(start, int):
+        start = (start,) * rank  # broadcast to all dimensions
+    elif isinstance(start, (list, tuple)):
+        if len(start) != rank:
+            raise ValueError("Start list length does not match dataset rank")
+    else:
+        raise TypeError("Start value must be an int or a list/tuple of ints")
+
+    if isinstance(stop, int):
+        stop = (stop,) * rank  # broadcast to all dimensions
+    elif isinstance(stop, (list, tuple)):
+        if len(stop) != rank:
+            raise ValueError("Stop list length does not match dataset rank")
+    else:
+        raise TypeError("Stop value must be an int or a list/tuple of ints")
+
+    if isinstance(step, int):
+        step = (step,) * rank  # broadcast to all dimensions
+    elif isinstance(step, (list, tuple)):
+        if len(step) != rank:
+            raise ValueError("Step list length does not match dataset rank")
+    else:
+        raise TypeError("Step value must be an int or a list/tuple of ints")
+
+    for idx in range(rank):
+        s = slice(start[idx], stop[idx], step[idx])
+        slices.append(s)
+
+    return slices
+
+
 class Selection(object):
 
     """

test/unit/selection_test.py

@@ -60,6 +60,10 @@ def testSelectAll(self):
         self.assertEqual(sel, sel2)
         self.assertEqual(sel.query_string, query_string)
 
+        # create the same selection from a dict
+        sel3 = selections.select(shape, {})
+        self.assertEqual(sel, sel3)
+
     def testSelectAll2D(self):
         shape = (4, 5)
         sel = selections.select(shape, ...)
@@ -80,6 +84,10 @@ def testSelectAll2D(self):
         self.assertEqual(sel, sel2)
         self.assertEqual(sel.query_string, query_string)
 
+        # create the same selection from a dict
+        sel3 = selections.select(shape, {})
+        self.assertEqual(sel, sel3)
+
     def testSlice1D(self):
         shape = (10,)
         sel = selections.select(shape, slice(2, 7))
@@ -102,6 +110,10 @@ def testSlice1D(self):
         self.assertEqual(sel, sel2)
         self.assertEqual(sel.query_string, query_string)
 
+        # create the same selection from a dict
+        sel3 = selections.select(shape, {"start": 2, "stop": 7})
+        self.assertEqual(sel, sel3)
+
     def testSliceWithStep(self):
         shape = (10,)
         sel = selections.select(shape, slice(0, 10, 2))
@@ -124,6 +136,10 @@ def testSliceWithStep(self):
         self.assertEqual(sel, sel2)
         self.assertEqual(sel.query_string, query_string)
 
+        # create the same selection from a dict
+        sel3 = selections.select(shape, {"start": 0, "stop": 10, "step": 2})
+        self.assertEqual(sel, sel3)
+
     def testSlice2D(self):
         shape = (8, 10)
         sel = selections.select(shape, (slice(1, 4), slice(2, 9)))
@@ -146,6 +162,10 @@ def testSlice2D(self):
         self.assertEqual(sel, sel2)
         self.assertEqual(sel.query_string, query_string)
 
+        # create the same selection from a dict
+        sel3 = selections.select(shape, {"start": [1, 2], "stop": [4, 9]})
+        self.assertEqual(sel, sel3)
+
     def testBroadcast1D(self):
         shape = (10,)
         sel = selections.select(shape, ...)
@@ -388,6 +408,7 @@ def testFancyCoord(self):
         self.assertEqual(len(slices), 2)
         self.assertEqual(slices[0], slice(0, 5, 1))
         self.assertEqual(slices[1], [3, 7])
+
         # create the same selection from a string
         query_string = "[0:5,[3,7]]"
         sel2 = selections.select(shape, query_string)