fix naming func

flint/naming.py

@@ -7,7 +7,7 @@
 import re
 from datetime import datetime
 from pathlib import Path
-from typing import Any, NamedTuple
+from typing import Any, NamedTuple, overload
 
 from prefect import task
 
@@ -16,22 +16,33 @@
 from flint.options import MS
 
 
-def rename_linear_to_stokes(
-    linear_name: Path | str,
+def _rename_linear_to_stokes(
+    linear_name_str: str,
     stokes: str,
-) -> Path | str:
+) -> str:
     if stokes.lower() not in ("q", "u"):
         raise NameError(f"Stokes {stokes=} is not linear!")
-
     pattern = r"\.qu\."  # Regex pattern to replace
+    stokes_name = re.sub(pattern, stokes, linear_name_str)
+    return stokes_name
+
+
+@overload
+def rename_linear_to_stokes(linear_name: Path, stokes: str) -> Path: ...
+
 
+@overload
+def rename_linear_to_stokes(linear_name: str, stokes: str) -> str: ...
+
+
+def rename_linear_to_stokes(
+    linear_name: Path | str,
+    stokes: str,
+) -> Path | str:
     if isinstance(linear_name, Path):
-        name_str = linear_name.as_posix()
-        stokes_name = re.sub(pattern, stokes, name_str)
-        return Path(stokes_name)
-    else:
-        stokes_name = re.sub(pattern, stokes, linear_name)
-        return stokes_name
+        return Path(_rename_linear_to_stokes(linear_name.as_posix(), stokes))
+
+    return _rename_linear_to_stokes(linear_name, stokes)
 
 
 task_rename_linear_to_stokes = task(rename_linear_to_stokes)

flint/prefect/flows/polarisation_pipeline.py

@@ -160,7 +160,10 @@ def process_science_fields_pol(
                         if polarisation == "linear":
                             # Get single Stokes prefix - the original prefix is the linear prefix
                             # i.e. `.qu.` -> `.q.` or `.u.` depending on the stokes
-                            prefix = task_rename_linear_to_stokes.submit(prefix)
+                            prefix = task_rename_linear_to_stokes.submit(
+                                linear_name=prefix,
+                                stokes=stokes,
+                            )
                         cube_path = combine_images_to_cube.submit(
                             images=channel_image_list,
                             prefix=prefix,