diff --git a/src/spyglass/utils/dj_helper_fn.py b/src/spyglass/utils/dj_helper_fn.py index 42cf67ba0..9142eb26a 100644 --- a/src/spyglass/utils/dj_helper_fn.py +++ b/src/spyglass/utils/dj_helper_fn.py @@ -237,7 +237,7 @@ def get_nwb_table(query_expression, tbl, attr_name, *attrs, **kwargs): # TODO: check that the query_expression restricts tbl - CBroz nwb_files = ( query_expression.join( - tbl.proj(nwb2load_filepath=attr_name), log_export=False + tbl.proj(nwb2load_filepath=attr_name), ) ).fetch(file_name_str) @@ -274,6 +274,8 @@ def fetch_nwb(query_expression, nwb_master, *attrs, **kwargs): """ kwargs["as_dict"] = True # force return as dictionary + close_file = kwargs.pop("close_file", True) + tbl, attr_name = nwb_master if "analysis" in attr_name: file_name_attr = "analysis_file_name" @@ -294,7 +296,7 @@ def fetch_nwb(query_expression, nwb_master, *attrs, **kwargs): get_nwb_file(file_path) query_table = query_expression.join( - tbl.proj(nwb2load_filepath=attr_name), log_export=False + tbl.proj(nwb2load_filepath=attr_name), ) rec_dicts = query_table.fetch(*attrs, **kwargs) # get filepath for each. Use datajoint for checksum if local @@ -318,7 +320,7 @@ def fetch_nwb(query_expression, nwb_master, *attrs, **kwargs): ret = [] for rec_dict in rec_dicts: - nwbf = get_nwb_file(rec_dict.pop("nwb2load_filepath")) + nwbf, io = get_nwb_file(rec_dict.pop("nwb2load_filepath"), return_io=True) # for each attr that contains substring 'object_id', store key-value: attr name to NWB object # remove '_object_id' from attr name nwb_objs = { @@ -329,6 +331,8 @@ def fetch_nwb(query_expression, nwb_master, *attrs, **kwargs): if "object_id" in id_attr and rec_dict[id_attr] != "" } ret.append({**rec_dict, **nwb_objs}) + if close_file: + io.close() return ret diff --git a/src/spyglass/utils/nwb_helper_fn.py b/src/spyglass/utils/nwb_helper_fn.py index 82af8a626..3fe0bbd0c 100644 --- a/src/spyglass/utils/nwb_helper_fn.py +++ b/src/spyglass/utils/nwb_helper_fn.py @@ -22,7 +22,7 @@ invalid_electrode_index = 99999999 -def get_nwb_file(nwb_file_path): +def get_nwb_file(nwb_file_path, **kwargs): """Return an NWBFile object with the given file path in read mode. If the file is not found locally, this will check if it has been shared @@ -45,7 +45,7 @@ def get_nwb_file(nwb_file_path): nwb_file_path = Nwbfile.get_abs_path(nwb_file_path) - _, nwbfile = __open_nwb_files.get(nwb_file_path, (None, None)) + io, nwbfile = __open_nwb_files.get(nwb_file_path, (None, None)) if nwbfile is None: # check to see if the file exists @@ -96,6 +96,11 @@ def get_nwb_file(nwb_file_path): nwbfile = io.read() __open_nwb_files[nwb_file_path] = (io, nwbfile) + __open_nwb_files.clear() + return_io = kwargs.pop("return_io", False) + if return_io: + return nwbfile, io + return nwbfile