astrodatadownloader: show later/better matching stage for fetching which files to download

main/astrodatadownloader/astrodatadownloader.py

@@ -57,7 +57,19 @@ def parse_manifest(manifest):
 
     return results
 
-def getStarObservations(starName: str, sources: list[ObservationSource], sequences: list[list[int]] = []) -> str:
+def fitsFilenameFilterFunc(table, key_colnames):
+    if(str(table["productFilename"]).lower().endswith(".fits") or
+       str(table["productFilename"]).lower().endswith(".fit")):
+        return True
+    return False
+
+def getStarObservations(starName: str, keplerCadences, k2Cadences,
+                        sources: list[ObservationSource], sequences: list[list[int]] = []) -> str:
+    if(len(keplerCadences) != 2):
+        raise Exception("keplerCadences needs to have length 2")
+    if(len(k2Cadences) != 2):
+        raise Exception("k2Cadences need to have length 2")
+
     obs = Observations.query_object(objectname=starName, radius="0 deg")
 
     obsWantedFilter = obs["dataproduct_type"] == "timeseries"
@@ -79,21 +91,7 @@ def getStarObservations(starName: str, sources: list[ObservationSource], sequenc
 
     obsWantedFilter &= obsSourceFilter
 
-    return obs[obsWantedFilter]
-
-def fitsFilenameFilterFunc(table, key_colnames):
-    if(str(table["productFilename"]).lower().endswith(".fits") or
-       str(table["productFilename"]).lower().endswith(".fit")):
-        return True
-    return False
-
-def downloadStarProducts(obs_wanted, keplerCadences, k2Cadences):
-    if(len(keplerCadences) != 2):
-        raise Exception("keplerCadences needs to have length 2")
-    if(len(k2Cadences) != 2):
-        raise Exception("k2Cadences need to have length 2")
-
-    dataProducts = Observations.get_product_list(obs_wanted)
+    dataProducts = Observations.get_product_list(obs[obsWantedFilter])
 
     # LC:  TESS lightcurve
     # SLC: Kepler short cadence lightcurve
@@ -123,5 +121,15 @@ def downloadStarProducts(obs_wanted, keplerCadences, k2Cadences):
     productsWanted = productsWanted.group_by("productFilename")
     productsWanted = productsWanted.groups.filter(fitsFilenameFilterFunc)
 
+    newCol = []
+    for pwCol in productsWanted:
+        for owCol in obs[obsWantedFilter]:
+            if(pwCol["obsID"] == owCol["obsid"]):
+                newCol.append(owCol["sequence_number"])
+
+    productsWanted.add_column(newCol, name="sequence_number")
+    return productsWanted
+
+def downloadStarProducts(productsWanted):
     filenames = Observations.download_products(productsWanted)
     return parse_manifest(filenames)