From 1d407ac1b3d8c14dbd8b44887c45e24086606772 Mon Sep 17 00:00:00 2001
From: Samuel Bray <sam.bray@ucsf.edu>
Date: Thu, 21 Dec 2023 09:13:14 -0800
Subject: [PATCH] Helper function for adding data to kachery (#729)

* add helper function for adding data to kachery

* Use dj config default for default kachery zone

* insert to kacher selection in single call
---
 src/spyglass/sharing/__init__.py        |  1 +
 src/spyglass/sharing/sharing_kachery.py | 41 +++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/src/spyglass/sharing/__init__.py b/src/spyglass/sharing/__init__.py
index cdb5715a1..4a34a113b 100644
--- a/src/spyglass/sharing/__init__.py
+++ b/src/spyglass/sharing/__init__.py
@@ -2,4 +2,5 @@
     AnalysisNwbfileKachery,
     AnalysisNwbfileKacherySelection,
     KacheryZone,
+    share_data_to_kachery,
 )
diff --git a/src/spyglass/sharing/sharing_kachery.py b/src/spyglass/sharing/sharing_kachery.py
index 43e7eed71..0d7405b46 100644
--- a/src/spyglass/sharing/sharing_kachery.py
+++ b/src/spyglass/sharing/sharing_kachery.py
@@ -5,6 +5,7 @@
 from datajoint.errors import DataJointError
 
 from spyglass.utils.dj_mixin import SpyglassMixin
+from spyglass.settings import config
 
 from ..common.common_lab import Lab  # noqa: F401
 from ..common.common_nwbfile import AnalysisNwbfile
@@ -215,3 +216,43 @@ def download_file(analysis_file_name: str) -> bool:
             raise Exception(f"{analysis_file_name} cannot be downloaded")
 
         return True
+
+
+def share_data_to_kachery(
+    restriction={},
+    table_list=[],
+    zone_name=None,
+):
+    """Share data to kachery
+
+    Parameters
+    ----------
+    restriction : dict, optional
+        restriction to select what data should be shared from table, by default {}
+    table_list : list, optional
+        List of tables to share data from, by default []
+    zone_name : str, optional
+        What kachery zone to share the data to, by default zone in spyglass.settings.config,
+        which looks for `KACHERY_ZONE` environmental variable, but defaults to
+        'franklab.default'
+
+    Raises
+    ------
+    ValueError
+        Does not allow sharing of all data in table
+    """
+    if not zone_name:
+        zone_name = config["KACHERY_ZONE"]
+    kachery_selection_key = {"kachery_zone_name": zone_name}
+    if not restriction:
+        raise ValueError("Must provide a restriction to the table")
+    selection_inserts = []
+    for table in table_list:
+        analysis_file_list = (table & restriction).fetch("analysis_file_name")
+        for file in analysis_file_list:  # Add all analysis to shared list
+            kachery_selection_key["analysis_file_name"] = file
+            selection_inserts.append(kachery_selection_key)
+    AnalysisNwbfileKacherySelection.insert(
+        selection_inserts, skip_duplicates=True
+    )
+    AnalysisNwbfileKachery.populate()