retrieve bursts by ID (#37)

* retrieve bursts by ID and  print found and not found burst IDs

Author: LiangJYu

src/s1reader/s1_reader.py

@@ -425,7 +425,8 @@ def _is_zip_annotation_xml(path: str, id_str: str) -> bool:
         return True
     return False
 
-def load_bursts(path: str, orbit_path: str, swath_num: int, pol: str = 'vv'):
+def load_bursts(path: str, orbit_path: str, swath_num: int, pol: str='vv',
+                burst_ids: list[str]=None):
     '''Find bursts in a Sentinel-1 zip file or a SAFE structured directory.
 
     Parameters:
@@ -438,16 +439,27 @@ def load_bursts(path: str, orbit_path: str, swath_num: int, pol: str = 'vv'):
         Integer of subswath of desired burst. {1, 2, 3}
     pol : str
         Polarization of desired burst. {hh, vv, hv, vh}
+    burst_ids : list[str]
+        List of burst IDs for which their Sentinel1BurstSlc objects will be
+        returned. Default of None returns all bursts. Empty list returned if
+        none of the burst IDs are found. If not all burst IDs are found, a list
+        containing found bursts will be returned.
 
     Returns:
     --------
     bursts : list
         List of Sentinel1BurstSlc objects found in annotation XML.
     '''
-
     if swath_num < 1 or swath_num > 3:
         raise ValueError("swath_num not <1 or >3")
 
+    if burst_ids is None:
+        burst_ids = []
+
+    # ensure burst IDs is a list
+    if not isinstance(burst_ids, list):
+        burst_ids = [burst_ids]
+
     # lower case polarity to be consistent with file naming convention
     pol = pol.lower()
     pols = ['vv', 'vh', 'hh', 'hv']
@@ -459,12 +471,30 @@ def load_bursts(path: str, orbit_path: str, swath_num: int, pol: str = 'vv'):
     if not os.path.exists(path):
         raise FileNotFoundError(f'{path} not found')
     elif os.path.isdir(path):
-        return _burst_from_safe_dir(path, id_str, orbit_path)
+        bursts = _burst_from_safe_dir(path, id_str, orbit_path)
     elif os.path.isfile(path):
-        return _burst_from_zip(path, id_str, orbit_path)
+        bursts = _burst_from_zip(path, id_str, orbit_path)
     else:
         raise ValueError(f'{path} is unsupported')
 
+    if burst_ids:
+        bursts = [b for b in bursts if b.burst_id in burst_ids]
+
+        burst_ids_found = set([b.burst_id for b in bursts])
+
+        if not burst_ids_found:
+            warnings.warn("None of provided bursts IDs found")
+
+        set_burst_ids = set(burst_ids)
+        if burst_ids_found != set_burst_ids:
+            diff = set_burst_ids.difference(burst_ids_found)
+            warn_str = f'Not all burst IDs found. Not found: {diff}. '
+            warn_str += f'Found: {burst_ids_found}'
+            warnings.warn(warn_str)
+
+    return bursts
+
+
 def _burst_from_zip(zip_path: str, id_str: str, orbit_path: str):
     '''Find bursts in a Sentinel-1 zip file.