Formatting

Author: Alec Thomson

arrakis/frion.py

@@ -6,17 +6,16 @@
 import argparse
 import logging
 import os
+from collections.abc import Callable
 from pathlib import Path
 from pprint import pformat
-from typing import Callable
 from typing import NamedTuple as Struct
 from urllib.error import URLError
 
 import astropy.units as u
 import numpy as np
 import pymongo
 from astropy.time import Time, TimeDelta
-from FRion import correct, predict
 from prefect import flow, get_run_logger, task
 from tqdm.auto import tqdm
 
@@ -29,6 +28,7 @@
 )
 from arrakis.utils.fitsutils import getfreq
 from arrakis.utils.pipeline import generic_parser, logo_str, workdir_arg_parser
+from FRion import correct, predict
 
 logger.setLevel(logging.INFO)
 TQDM_OUT = TqdmToLogger(logger, level=logging.INFO)

arrakis/linmos.py

@@ -12,6 +12,7 @@
 from pathlib import Path
 from pprint import pformat
 from typing import NamedTuple as Struct
+from typing import overload
 
 import astropy.units as u
 import numpy as np
@@ -216,9 +217,21 @@ def genparset(
     return parset_file
 
 
+@overload
+def linmos(
+    parset: None, fieldname: str, image: str, holofile: Path | None = None
+) -> None: ...
+
+
+@overload
+def linmos(
+    parset: str, fieldname: str, image: str, holofile: Path | None = None
+) -> pymongo.UpdateOne: ...
+
+
 @task(name="Run linmos")
 def linmos(
-    parset: str | None, fieldname: str, image: str, holofile: Path
+    parset: str | None, fieldname: str, image: str, holofile: Path | None = None
 ) -> pymongo.UpdateOne | None:
     """Run linmos
 
@@ -248,12 +261,15 @@ def linmos(
     log_file = parset.replace(".in", ".log")
     linmos_command = shlex.split(f"linmos -c {parset}")
 
-    holo_folder = holofile.parent
+    bind_dir_str = f"{rootdir}:{rootdir}"
+    if holofile is not None:
+        holo_folder = holofile.parent
+        bind_dir_str += f",{holo_folder}:{holo_folder}"
 
     output = sclient.execute(
         image=image,
         command=linmos_command,
-        bind=f"{rootdir}:{rootdir},{holo_folder}:{holo_folder}",
+        bind=bind_dir_str,
         return_result=True,
         quiet=False,
         stream=True,

arrakis/makecat.py

@@ -8,9 +8,10 @@
 import os
 import time
 import warnings
+from collections.abc import Callable
 from pathlib import Path
 from pprint import pformat
-from typing import Callable, NamedTuple
+from typing import NamedTuple
 
 import astropy.units as u
 import dask.dataframe as dd