Add option to select datatype during query

atdb_csv_gen/atdb_csv_gen/csv_gen.py

@@ -13,6 +13,10 @@ from sqlalchemy.engine.base import Connection, Engine
 from sqlalchemy.sql import text
 
 logger = logging.getLogger(__name__)
+_EXPECTED_STR_PER_DATATYPE = {
+    'MS': '.MS_',
+    'BF': '_bf.'
+}
 
 
 def parse_args() -> Namespace:
@@ -32,6 +36,7 @@ def parse_args() -> Namespace:
     parser.add_argument(
         "--save-missing", type=str, help="Directory to store csv with missing files"
     )
+    parser.add_argument("--datatype", help="Select a specific datatype", choices=_EXPECTED_STR_PER_DATATYPE.keys())
     parser.add_argument("-v", action="store_true", help="Verbose logging")
     parser.add_argument("-q", action="store_true", help="Quiet logging (only errors)")
     return parser.parse_args()
@@ -156,11 +161,20 @@ def match_file_lists(
     return reg_data.intersection(cra_data)
 
 
+def filter_datatype(size_path_list, datatype):
+    if datatype is None:
+        return size_path_list
+    else:
+        expected_string = _EXPECTED_STR_PER_DATATYPE[datatype]
+        return [item for item in size_path_list if expected_string in item]
+
+
 def gen_csv(
         obs_id: str,
         o_file: Optional[str] = None,
         save_missing: Optional[str] = None,
         config_file: Optional[str] = None,
+        datatype: Optional[str] = None,
 ):
     """Generate CSV file for ATDB
 
@@ -180,7 +194,9 @@ def gen_csv(
 
     intersection = match_file_lists(registered, crawled, save_missing)
 
-    write_output("\n".join(intersection), o_file)
+    filtered_intersection = filter_datatype(intersection, datatype)
+
+    write_output("\n".join(filtered_intersection), o_file)
 
 
 def main():
@@ -198,6 +214,7 @@ def main():
         o_file=args.o,
         save_missing=args.save_missing,
         config_file=args.c,
+        datatype=args.datatype
     )