mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-09-15 20:57:56 +08:00
289 lines
11 KiB
Python
289 lines
11 KiB
Python
import base64
|
|
import os
|
|
import re
|
|
import threading
|
|
from datetime import datetime
|
|
from hashlib import blake2b
|
|
|
|
import pandas as pd
|
|
from flask import session
|
|
from fsspec import AbstractFileSystem
|
|
|
|
from server import __version__ as cellxgene_version
|
|
from server.app.session import get_user_id
|
|
from server.common.annotations.annotations import Annotations
|
|
from server.common.genesets import read_gene_sets_tidycsv
|
|
from server.common.errors import AnnotationsError, ObsoleteRequest
|
|
from server.common.utils.data_locator import DataLocator
|
|
|
|
|
|
class AnnotationsLocalFile(Annotations):
|
|
CXG_ANNO_COLLECTION = "cxg_anno_collection"
|
|
|
|
def __init__(self, config, output_dir, label_output_file, gene_sets_output_file):
|
|
super().__init__(config)
|
|
self.output_dir = output_dir
|
|
self.label_output_file = label_output_file
|
|
self.gene_sets_output_file = gene_sets_output_file
|
|
# lock used to protect label file write ops
|
|
self.label_lock = threading.RLock()
|
|
self.gene_sets_lock = threading.RLock()
|
|
|
|
# cache the most recent cell labels/annotations.
|
|
self.last_label_fname = None
|
|
self.last_labels = None
|
|
|
|
# cache the most recent gene sets.
|
|
self.last_geneset_fname = None
|
|
self.last_geneset = None
|
|
|
|
# txn ID - used to de-dup geneset writes
|
|
self.last_geneset_tid = 0
|
|
|
|
def is_safe_collection_name(self, name):
|
|
"""
|
|
return true if this is a safe collection name
|
|
this is ultra conservative. If we want to allow full legal file name syntax,
|
|
we could look at modules like `pathvalidate`
|
|
"""
|
|
if name is None:
|
|
return False
|
|
return re.match(r"^[\w\-]+$", name) is not None
|
|
|
|
def set_collection(self, name):
|
|
session[self.CXG_ANNO_COLLECTION] = name
|
|
session.permanent = True
|
|
|
|
def get_collection(self):
|
|
if session is None:
|
|
return None
|
|
return session.get(self.CXG_ANNO_COLLECTION)
|
|
|
|
def read_labels(self, data_adaptor):
|
|
self.check_user_annotations_enabled() # raises
|
|
|
|
fname = self._get_celllabels_filename(data_adaptor)
|
|
empty_labels = pd.DataFrame()
|
|
if fname is None:
|
|
return empty_labels
|
|
|
|
with self.label_lock:
|
|
locator = DataLocator(fname)
|
|
if not locator.exists() or locator.size() == 0:
|
|
return empty_labels
|
|
|
|
# return the cached labels if possible
|
|
if fname == self.last_label_fname:
|
|
return self.last_labels
|
|
|
|
# otherwise, read labels from file
|
|
with locator.open() as f:
|
|
labels = pd.read_csv(f, dtype="category", index_col=0, header=0, comment="#", keep_default_na=False)
|
|
|
|
# update the cache
|
|
self.last_label_fname = fname
|
|
self.last_labels = labels
|
|
return labels
|
|
|
|
def write_labels(self, df, data_adaptor):
|
|
self.check_user_annotations_enabled() # raises
|
|
|
|
# update our internal state and save it. Multi-threading often enabled,
|
|
# so treat this as a critical section.
|
|
with self.label_lock:
|
|
lastmod = data_adaptor.get_last_mod_time()
|
|
lastmodstr = "'unknown'" if lastmod is None else lastmod.isoformat(timespec="seconds")
|
|
header = (
|
|
f"# Annotations generated on {datetime.now().isoformat(timespec='seconds')} "
|
|
f"using cellxgene version {cellxgene_version}\n"
|
|
f"# Input data file was {data_adaptor.get_location()}, "
|
|
f"which was last modified on {lastmodstr}\n"
|
|
)
|
|
|
|
fname = self._get_celllabels_filename(data_adaptor)
|
|
self._backup(fname)
|
|
locator = DataLocator(fname)
|
|
with locator.open("w") as f:
|
|
if not df.empty:
|
|
if header is not None:
|
|
f.write(header)
|
|
df.to_csv(f)
|
|
|
|
# update the cache
|
|
self.last_label_fname = fname
|
|
self.last_labels = df
|
|
|
|
def read_gene_sets(self, data_adaptor, context=None):
|
|
fname = self._get_genesets_filename(data_adaptor)
|
|
empty_gene_sets = {}
|
|
|
|
with self.gene_sets_lock:
|
|
tid = self.last_geneset_tid # inside the critical section
|
|
if fname is None:
|
|
return (empty_gene_sets, tid)
|
|
|
|
locator = DataLocator(fname)
|
|
if not locator.exists() or locator.size() == 0:
|
|
return (empty_gene_sets, tid)
|
|
|
|
# return the cached genesets if possible, otherwise read from file and validate them
|
|
if fname == self.last_geneset_fname:
|
|
return (self.last_geneset, tid)
|
|
|
|
# read
|
|
gene_sets = read_gene_sets_tidycsv(locator, context)
|
|
|
|
# validate
|
|
gene_sets = data_adaptor.check_new_gene_sets(gene_sets, context)
|
|
|
|
# update cache
|
|
self.last_geneset_fname = fname
|
|
self.last_geneset = gene_sets
|
|
|
|
return (gene_sets, tid)
|
|
|
|
def write_gene_sets(self, gene_sets, tid, data_adaptor):
|
|
self.check_gene_sets_save_enabled() # raises
|
|
|
|
if type(tid) is not int or tid < 0:
|
|
raise ValueError("tid must be a positive integer")
|
|
|
|
# may raise
|
|
gene_sets = data_adaptor.check_new_gene_sets(gene_sets)
|
|
|
|
with self.gene_sets_lock:
|
|
# skip if the request is stale
|
|
if tid is not None:
|
|
if tid <= self.last_geneset_tid:
|
|
raise ObsoleteRequest("TID is stale.")
|
|
self.last_geneset_tid = tid
|
|
|
|
lastmod = data_adaptor.get_last_mod_time()
|
|
lastmodstr = "'unknown'" if lastmod is None else lastmod.isoformat(timespec="seconds")
|
|
header = (
|
|
f"# Gene set generated on {datetime.now().isoformat(timespec='seconds')} "
|
|
f"using cellxgene version {cellxgene_version}\n"
|
|
f"# Input data file was {data_adaptor.get_location()}, "
|
|
f"which was last modified on {lastmodstr}\n"
|
|
)
|
|
|
|
fname = self._get_genesets_filename(data_adaptor)
|
|
self._backup(fname)
|
|
locator = DataLocator(fname)
|
|
with locator.open("w", newline="") as f:
|
|
f.write(header + self.gene_sets_to_csv(gene_sets))
|
|
|
|
# update the cache
|
|
self.last_geneset_fname = fname
|
|
self.last_geneset = gene_sets if isinstance(gene_sets, dict) else {g["geneset_name"]: g for g in gene_sets}
|
|
|
|
def _get_userdata_idhash(self, data_adaptor):
|
|
"""
|
|
Return a short hash that weakly identifies the user and dataset.
|
|
Used to create safe annotations output file names.
|
|
"""
|
|
uid = get_user_id(session)
|
|
id = (uid + data_adaptor.get_location()).encode()
|
|
idhash = base64.b32encode(blake2b(id, digest_size=5).digest()).decode("utf-8")
|
|
return idhash
|
|
|
|
def _get_output_dir(self):
|
|
if self.output_dir:
|
|
return self.output_dir
|
|
|
|
output_file = self.label_output_file or self.gene_sets_output_file
|
|
if output_file:
|
|
return os.path.dirname(DataLocator(output_file).abspath())
|
|
|
|
return os.getcwd()
|
|
|
|
def _get_celllabels_filename(self, data_adaptor):
|
|
"""return the current annotation file name"""
|
|
if self.label_output_file:
|
|
return self.label_output_file
|
|
|
|
return self._get_filename(data_adaptor, "cell-labels")
|
|
|
|
def _get_genesets_filename(self, data_adaptor):
|
|
"""return the current gene sets file name"""
|
|
if self.gene_sets_output_file:
|
|
return self.gene_sets_output_file
|
|
|
|
return self._get_filename(data_adaptor, "gene-sets")
|
|
|
|
def _get_filename(self, data_adaptor, anno_name):
|
|
# we need to generate a file name, which we can only do if we have a UID and collection name
|
|
if session is None:
|
|
raise AnnotationsError("unable to determine file name for annotations")
|
|
|
|
collection = self.get_collection()
|
|
if collection is None:
|
|
return None
|
|
|
|
if data_adaptor is None:
|
|
raise AnnotationsError("unable to determine file name for annotations")
|
|
|
|
idhash = self._get_userdata_idhash(data_adaptor)
|
|
return os.path.join(self._get_output_dir(), f"{collection}-{anno_name}-{idhash}.csv")
|
|
|
|
def _backup(self, fname, max_backups=9):
|
|
"""
|
|
save N backups of file to backup_dir.
|
|
1. fname -> backup_dir/fname-TIME
|
|
2. delete excess files in backup_dir
|
|
"""
|
|
locator = DataLocator(fname)
|
|
fs: AbstractFileSystem = locator.fs # Handle to underlying fsspec file system
|
|
|
|
# Make sure there is work to do
|
|
if not locator.exists():
|
|
return
|
|
|
|
root, ext = os.path.splitext(locator.abspath())
|
|
backup_dir = f"{root}-backups"
|
|
|
|
# Ensure backup_dir exists
|
|
fs.mkdirs(backup_dir, exist_ok=True)
|
|
|
|
# Save current file to backup_dir
|
|
fname_base = os.path.basename(fname)
|
|
fname_base_root, fname_base_ext = os.path.splitext(fname_base)
|
|
# don't use ISO standard time format, as it contains characters illegal on some filesystems.
|
|
nowish = datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
|
|
backup_fname = os.path.join(backup_dir, f"{fname_base_root}-{nowish}{fname_base_ext}")
|
|
if fs.exists(backup_fname):
|
|
fs.delete(backup_fname)
|
|
fs.rename(fname, backup_fname)
|
|
|
|
# prune the backup_dir to max number of backup files, keeping the most recent backups
|
|
backup_path_prefix = DataLocator.strip_protocol(os.path.join(backup_dir, fname_base_root + "-"))
|
|
backups = list(filter(lambda s: s.startswith(backup_path_prefix), fs.ls(backup_dir)))
|
|
|
|
# sorting to drop the oldest
|
|
excess_backups = list(sorted(backups, reverse=True))[max_backups:]
|
|
for bu in excess_backups:
|
|
fs.delete(bu)
|
|
|
|
def update_parameters(self, parameters, data_adaptor):
|
|
params = {}
|
|
params["annotations"] = self.user_annotations_enabled()
|
|
params["annotations_genesets_readonly"] = not self.gene_sets_save_enabled()
|
|
params["annotations_genesets_name_is_read_only"] = self.gene_sets_output_file is not None
|
|
params["user_annotation_collection_name_enabled"] = True
|
|
|
|
if self.label_output_file is not None:
|
|
# user has hard-wired the name of the annotation cell label data collection
|
|
fname = os.path.basename(self.label_output_file)
|
|
collection_fname = os.path.splitext(fname)[0]
|
|
params["annotations-data-collection-is-read-only"] = True
|
|
params["annotations-data-collection-name"] = collection_fname
|
|
|
|
elif session is not None:
|
|
collection = self.get_collection()
|
|
params["annotations-data-collection-is-read-only"] = not self.user_annotations_enabled()
|
|
params["annotations-data-collection-name"] = collection
|
|
|
|
params["annotations-user-data-idhash"] = self._get_userdata_idhash(data_adaptor)
|
|
|
|
parameters.update(params)
|