GDAL_Metadata instances are now subscriptable.

geoarray/baseclasses.py

@@ -1487,7 +1487,7 @@ class GeoArray(object):
         # apply zslice to bandnames and metadata
         if zslicing:
             bNs_out = list(np.array(list(self.bandnames))[zslice])
-            _meta_out = self.metadata.get_band_subset(bandslice=zslice)
+            _meta_out = self.metadata.get_subset(bands2extract=zslice)
         else:
             bNs_out = list(self.bandnames)
             _meta_out = self.meta

geoarray/metadata.py

@@ -3,6 +3,7 @@
 import os
 from pprint import pformat
 from copy import deepcopy
+from typing import Union  # noqa F401  # flake8 issue
 
 from geopandas import GeoDataFrame, GeoSeries
 import numpy as np
@@ -204,13 +205,60 @@ class GDAL_Metadata(object):
         return dict(zip(self.all_meta.keys(),
                         [self._convert_param_to_ENVI_str(i) for i in self.all_meta.values()]))
 
-    def get_band_subset(self, bandslice=None):
-        bandslice = bandslice or slice(None)
+    def get_subset(self, bands2extract=None, keys2extract=None):
+        # type: (Union[slice, list, np.ndarray], Union[str, list]) -> 'GDAL_Metadata'
         meta_sub = deepcopy(self)
 
-        for k, v in meta_sub.band_meta.items():
-            meta_sub.band_meta[k] = list(np.array(v)[bandslice])
+        # subset bands
+        if bands2extract is not None:
+            if isinstance(bands2extract, list):
+                bands2extract = np.array(bands2extract)
+            elif isinstance(bands2extract, (np.ndarray, slice)):
+                pass  # all fine
+            else:
+                raise TypeError(bands2extract)
+
+            for k, v in meta_sub.band_meta.items():
+                meta_sub.band_meta[k] = list(np.array(v)[bands2extract])
+
+            meta_sub.bands = len(list(range(*bands2extract.indices(bands2extract.stop)))) \
+                if isinstance(bands2extract, slice) else bands2extract.size
+
+        # subset metadata keys
+        if keys2extract:
+            keys2extract = [keys2extract] if isinstance(keys2extract, str) else keys2extract
+
+            # global_meta = meta_sub.global_meta.copy()
+            for k in meta_sub.global_meta.copy().keys():
+                if k not in keys2extract:
+                    del meta_sub.global_meta[k]
 
-        meta_sub.bands = len(list(range(*bandslice.indices(bandslice.stop))))
+            for k in meta_sub.band_meta.copy().keys():
+                if k not in keys2extract:
+                    del meta_sub.band_meta[k]
+
+            if not meta_sub.all_meta:
+                raise ValueError(keys2extract, 'The given metadata keys do not exist.')
 
         return meta_sub
+
+    def __getitem__(self, given):
+        if isinstance(given, int):
+            return self.get_subset(bands2extract=slice(given, given + 1))
+        elif isinstance(given, slice):
+            return self.get_subset(bands2extract=given)
+        elif isinstance(given, str):
+            return self.get_subset(keys2extract=given)
+        elif isinstance(given, list):
+            if isinstance(given[0], str):
+                return self.get_subset(keys2extract=given)
+            elif isinstance(given[0], int):
+                return self.get_subset(bands2extract=given)
+            else:
+                raise TypeError(given, 'Given list must contain string or integer items.')
+        elif isinstance(given, np.ndarray):
+            if given.ndim != 1:
+                raise TypeError(given, 'Given numpy array must be one-dimensional.')
+            return self.get_subset(bands2extract=given)
+        else:
+            raise TypeError(given)