fix setting a partial stack of bands and add test

backport to 0.7.x
fortyninemaps · Oct 8, 2016 · 9e5188a · 9e5188a
1 parent 9f05696
commit 9e5188a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/karta/raster/band.py b/karta/raster/band.py
@@ -86,8 +86,8 @@ def __setitem__(self, key, value):
                 sb = slice(None, None, None) if len(key) == 2 else key[2]
                 if isinstance(sb, slice):
                     if len(value.shape) == 3:
-                        for i in range(len(self.bands)):
-                            self.bands[i][sr,sc] = value[:,:,i]
+                        for ival, iband in enumerate(range(*sb.indices(len(self.bands)))):
+                            self.bands[iband][sr,sc] = value[:,:,ival]
                     else:
                         for b in self.bands[sb]:
                             b[sr,sc] = value

diff --git a/tests/band_tests.py b/tests/band_tests.py
@@ -132,6 +132,11 @@ def test_set_multibanded(self):
         self.assertTrue(np.all(bands[0][:,:] == 1.0))
         self.assertTrue(np.all(bands[1][:,:] == 2.0))
         self.assertTrue(np.all(bands[2][:,:] == 2.0))
+
+        indexer[:,:,1:] = np.dstack([2*values, 3*values])
+
+        self.assertTrue(np.all(bands[1][:,:] == 2.0))
+        self.assertTrue(np.all(bands[2][:,:] == 3.0))
         return
 
     def test_set_multibanded_masked(self):