From 8ba3e7892841b3bf6109640e96a316a658b58d49 Mon Sep 17 00:00:00 2001
From: bjhardcastle <ben.hardcastle@alleninstitute.org>
Date: Tue, 28 May 2024 13:24:44 -0700
Subject: [PATCH] Consolidate nwb writing methods and check file exists

---
 src/npc_sessions/sessions.py | 40 +++++++++++++-----------------------
 1 file changed, 14 insertions(+), 26 deletions(-)

diff --git a/src/npc_sessions/sessions.py b/src/npc_sessions/sessions.py
index 516bcf9..16c59da 100644
--- a/src/npc_sessions/sessions.py
+++ b/src/npc_sessions/sessions.py
@@ -348,8 +348,8 @@ def nwb(self) -> pynwb.NWBFile:
             units=self.units if self.is_sorted else None,
         )
 
-    def write_nwb_hdf5(
-        self, path: str | npc_io.PathLike | None = None, metadata_only: bool = False
+    def write_nwb(
+        self, path: str | npc_io.PathLike | None = None, metadata_only: bool = False, zarr=True, force=False,
     ) -> upath.UPath:
         """Write NWB file to disk - file path is normalized and returned"""
         if path is None:
@@ -358,32 +358,20 @@ def write_nwb_hdf5(
             path = npc_io.from_pathlike(path)
         path = path.with_stem(
             path.name.replace(".hdf5", "").replace(".nwb", "").replace(".zarr", "")
-        ).with_suffix(".nwb")
+        ).with_suffix(".nwb.zarr" if zarr else ".nwb")
+        if not force and path.exists() and (npc_io.get_size(path) // 1024) > 1:
+            raise FileExistsError(f"{path} already exists - use `force=True` to overwrite")
+        elif force and zarr and path.exists():
+            logger.warning(f"Overwriting zarr directories is not advised: remnants of previous data may remain.\nSuggest deleting {path} first.")
         nwb = self.nwb if not metadata_only else self.metadata
-        with pynwb.NWBHDF5IO(path.as_posix(), "w") as io:
-            io.write(nwb)
-        logger.info(
-            f"Saved NWB file to {path}: {npc_io.get_size(path) // 1024 ** 2} MB"
-        )
-        return path
-
-    def write_nwb_zarr(
-        self, path: str | npc_io.PathLike | None = None, metadata_only: bool = False
-    ) -> upath.UPath:
-        """Write NWB file to disk - file path is normalized and returned"""
-        if path is None:
-            path = npc_lims.get_nwb_path(self.id)
+        if zarr:
+            with hdmf_zarr.NWBZarrIO(path.as_posix(), "w") as io:
+                io.write(
+                    nwb, link_data=False
+                )  # link_data=False so that lazily-opened zarrays are read and copied into nwb (instead of being added as a link, which is currently broken)
         else:
-            path = npc_io.from_pathlike(path)
-        path = path.with_stem(
-            path.name.replace(".hdf5", "").replace(".nwb", "").replace(".zarr", "")
-        ).with_suffix(".nwb.zarr")
-        nwb = self.nwb if not metadata_only else self.metadata
-        with hdmf_zarr.NWBZarrIO(path.as_posix(), "w") as io:
-            io.write(
-                nwb, link_data=False
-            )  # link_data=False so that lazily-opened zarrays are read and copied into nwb (instead of being added as a link, which is currently broken)
-        # zarr "file" is a directory:
+            with pynwb.NWBHDF5IO(path.as_posix(), "w") as io:
+                io.write(nwb)
         logger.info(
             f"Saved NWB file to {path}: {npc_io.get_size(path) // 1024 ** 2} MB"
         )