indicate datetime type

punch-mission · Nov 26, 2024 · b469cac · b469cac
1 parent e3c1316
commit b469cac
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/punchpipe/flows/fcorona.py b/punchpipe/flows/fcorona.py
@@ -15,7 +15,8 @@
 
 
 @task
-def f_corona_background_query_ready_files(session, pipeline_config: dict, use_n: int = 250, reference_time=None):
+def f_corona_background_query_ready_files(session, pipeline_config: dict, use_n: int = 250,
+                                          reference_time: datetime =None):
     reference_time = reference_time or datetime.now()
     before = reference_time - timedelta(weeks=2)
     after = reference_time + timedelta(weeks=2)
@@ -40,7 +41,7 @@ def construct_f_corona_background_flow_info(level3_files: list[File],
                                             level3_f_model_file: File,
                                             pipeline_config: dict,
                                             session=None,
-                                            reference_time=None):
+                                            reference_time: datetime = None):
     flow_type = "construct_f_corona_background_process_flow"
     state = "planned"
     creation_time = datetime.now()
@@ -65,7 +66,8 @@ def construct_f_corona_background_flow_info(level3_files: list[File],
 
 
 @task
-def construct_f_corona_background_file_info(level2_files: t.List[File], pipeline_config: dict, reference_time=None) -> t.List[File]:
+def construct_f_corona_background_file_info(level2_files: t.List[File], pipeline_config: dict,
+                                            reference_time:datetime = None) -> t.List[File]:
     return [File(
                 level="3",
                 file_type="PF",
@@ -77,7 +79,7 @@ def construct_f_corona_background_file_info(level2_files: t.List[File], pipeline
             ),]
 
 @flow
-def f_corona_scheduler(pipeline_config_path=None, session=None, reference_time=None):
+def f_corona_scheduler(pipeline_config_path=None, session=None, reference_time: datetime = None):
     if not isinstance(reference_time, datetime):
         reference_time = parse_datetime_str(reference_time)
 

diff --git a/punchpipe/flows/starfield.py b/punchpipe/flows/starfield.py
@@ -13,7 +13,7 @@
 
 
 @task
-def starfield_background_query_ready_files(session, pipeline_config: dict, reference_time=None):
+def starfield_background_query_ready_files(session, pipeline_config: dict, reference_time: datetime =None):
     logger = get_run_logger()
     all_ready_files = (session.query(File)
                        .filter(File.state == "created")
@@ -31,7 +31,7 @@ def starfield_background_query_ready_files(session, pipeline_config: dict, refer
 def construct_starfield_background_flow_info(level3_fcorona_subtracted_files: list[File],
                                              level3_starfield_model_file: File,
                                              pipeline_config: dict,
-                                             session=None, reference_time=None):
+                                             session=None, reference_time: datetime =None):
     flow_type = "construct_starfield_background_process_flow"
     state = "planned"
     creation_time = datetime.now()