Support keys assign in GetterDataset

chainercv/chainer_experimental/datasets/sliceable/getter_dataset.py

@@ -49,13 +49,40 @@ class GetterDataset(SliceableDataset):
     def __init__(self):
         self._keys = []
         self._getters = []
+        self._return_tuple = True
 
     def __len__(self):
         raise NotImplementedError
 
     @property
     def keys(self):
-        return tuple(key for key, _, _ in self._keys)
+        if self._return_tuple:
+            return tuple(key for key, _, _ in self._keys)
+        else:
+            return self._keys[0][0]
+
+    @keys.setter
+    def keys(self, keys):
+        if isinstance(keys, (list, tuple)):
+            self._return_tuple = True
+        else:
+            keys, self._return_tuple = (keys,), False
+
+        new_keys = []
+        for key in keys:
+            if isinstance(key, int):
+                key_index = key
+                if key_index >= len(self._keys):
+                    raise IndexError('Invalid index of key')
+                if key_index < 0:
+                    key_index += len(self._keys)
+            else:
+                try:
+                    key_index = _as_tuple(self.keys).index(key)
+                except ValueError:
+                    raise KeyError('{} does not exists'.format(key))
+            new_keys.append(self._keys[key_index])
+        self._keys = new_keys
 
     def add_getter(self, keys, getter):
         """Register a getter function

tests/chainer_experimental_tests/datasets_tests/sliceable_tests/test_getter_dataset.py

@@ -45,5 +45,48 @@ def test_get_example_by_keys(self):
         self.assertEqual(example, ('item1(1)', 'item2(1)', 'item3(1)'))
         self.assertEqual(self.dataset.count, 2)
 
+    def test_set_keys_single_name(self):
+        self.dataset.keys = 'item0'
+        self.assertEqual(self.dataset.keys, 'item0')
+        self.assertEqual(self.dataset[1], 'item0(1)')
+
+    def test_set_keys_single_index(self):
+        self.dataset.keys = 0
+        self.assertEqual(self.dataset.keys, 'item0')
+        self.assertEqual(self.dataset[1], 'item0(1)')
+
+    def test_set_keys_single_tuple_name(self):
+        self.dataset.keys = ('item1',)
+        self.assertEqual(self.dataset.keys, ('item1',))
+        self.assertEqual(self.dataset[2], ('item1(2)',))
+
+    def test_set_keys_single_tuple_index(self):
+        self.dataset.keys = (1,)
+        self.assertEqual(self.dataset.keys, ('item1',))
+        self.assertEqual(self.dataset[2], ('item1(2)',))
+
+    def test_set_keys_multiple_name(self):
+        self.dataset.keys = ('item0', 'item2')
+        self.assertEqual(self.dataset.keys, ('item0', 'item2'))
+        self.assertEqual(self.dataset[3], ('item0(3)', 'item2(3)'))
+
+    def test_set_keys_multiple_index(self):
+        self.dataset.keys = (0, 2)
+        self.assertEqual(self.dataset.keys, ('item0', 'item2'))
+        self.assertEqual(self.dataset[3], ('item0(3)', 'item2(3)'))
+
+    def test_set_keys_multiple_mixed(self):
+        self.dataset.keys = ('item0', 2)
+        self.assertEqual(self.dataset.keys, ('item0', 'item2'))
+        self.assertEqual(self.dataset[3], ('item0(3)', 'item2(3)'))
+
+    def test_set_keys_invalid_name(self):
+        with self.assertRaises(KeyError):
+            self.dataset.keys = 'invalid'
+
+    def test_set_keys_invalid_index(self):
+        with self.assertRaises(IndexError):
+            self.dataset.keys = 4
+
 
 testing.run_module(__name__, __file__)