import logging
from datetime import datetime, timedelta
from mxcubecore.HardwareObjects.abstract.PyISPyBRestClient import (
PyISPyBRestClient,
PyISPyBUnsuccessfulResponse,
)
from mxcubecore.model.lims_session import LimsSessionManager, Proposal, Session
[docs]class PyISPyBDataAdapter:
"""Adapter to convert data from PyISPyB REST API."""
def __init__(
self,
client: PyISPyBRestClient,
beamline_name: str,
new_session_duration_days: int = 2,
):
self.client = client
self.beamline_name = beamline_name
self.new_session_duration_days = new_session_duration_days
self.logger = logging.getLogger("pyispyb_adapter")
def __is_time_between(
self, start_datetime: datetime, end_datetime: datetime
) -> bool:
"""Checks if the current time is between start and end."""
today = datetime.today() # noqa: DTZ002
try:
return start_datetime <= today <= end_datetime
except TypeError:
self.logger.exception(
"Invalid date format. start: %s, end: %s, now: %s",
start_datetime,
end_datetime,
today,
)
return False
def __to_proposal(self, proposal: dict[str, str]) -> Proposal:
"""
Converts proposal data received from PyISPyB REST API to a Proposal object.
"""
return Proposal(
code=proposal.get("proposalCode").upper(),
number=proposal.get("proposalNumber"),
proposal_id=str(proposal.get("proposalId")),
title=proposal.get("title"),
type=proposal.get("type", ""),
name=proposal.get("proposal"),
)
def __to_session(self, session: dict) -> Session:
"""Converts session data received from PyISPyB REST API to a Session object."""
proposal_name = session.get("proposal")
proposal_code = "".join([c for c in proposal_name if not c.isdigit()])
proposal_number = proposal_name[len(proposal_code) :]
title = session.get("title", "")
start_datetime = datetime.fromisoformat(session.get("startDate"))
end_datetime = datetime.fromisoformat(session.get("endDate"))
return Session(
code=proposal_code,
number=proposal_number,
proposal_name=proposal_name,
proposal_id=str(session.get("proposalId")),
session_id=str(session.get("sessionId")),
beamline_name=session.get("beamLineName", self.beamline_name),
title=title,
comments=session.get("comments"),
start_datetime=start_datetime,
end_datetime=end_datetime,
nb_shifts=str(session.get("nbShifts", "")),
scheduled=str(session.get("scheduled", "False")),
is_scheduled_time=self.__is_time_between(start_datetime, end_datetime),
is_scheduled_beamline=True, # MAX IV does not care about this value
)
@staticmethod
def __to_sample(sample_data: dict) -> dict:
"""
Convert sample data get from PyISPyB to internal
sample data structure.
Args:
sample_data (dict): The raw sample data from PyISPyB.
Returns:
dict: A dictionary conteining keys required by internal sample structure.
"""
container = sample_data["Container"]
crystal = sample_data["Crystal"]
return {
# experimentType
"sampleId": sample_data["blSampleId"],
"code": sample_data["code"],
"sampleName": sample_data["name"],
"sampleLocation": sample_data["location"],
"containerSampleChangerLocation": container["sampleChangerLocation"],
"containerCode": container["code"],
"proteinAcronym": crystal["Protein"]["acronym"],
"crystalId": crystal["crystalId"],
"crystalSpaceGroup": crystal["space_group"],
"cellA": crystal["cell_a"],
"cellAlpha": crystal["cell_alpha"],
"cellB": crystal["cell_b"],
"cellBeta": crystal["cell_beta"],
"cellC": crystal["cell_c"],
"cellGamma": crystal["cell_gamma"],
"diffractionPlan": sample_data["DiffractionPlan"]["diffractionPlanId"],
}
def __find_proposal_by_id(self, proposal_id: int) -> Proposal:
return self.__to_proposal(
self.client.get("proposals?proposalId=%s" % (proposal_id))[0]
)
[docs] def get_current_user_data(self) -> dict:
"""Fetches current user details."""
return self.client.get("user/current")
[docs] def get_proposals(self) -> list[Proposal]:
"""Returns proposals to which authenticated user has access."""
return [
self.__to_proposal(proposal) for proposal in self.client.get("proposals")
]
[docs] def find_proposal(self, code: str, number: str) -> Proposal:
"""Finds a proposal by its code and number."""
return self.__to_proposal(
self.client.get("proposals?search=%s%s" % (code, number))[0]
)
[docs] def get_sessions_by_code_and_number(
self, code: str, number: str, beamline: str
) -> LimsSessionManager:
"""Finds a session by its proposal code and number and beamline name."""
return LimsSessionManager(
sessions=[
self.__to_session(session)
for session in self.client.get(
"sessions?proposal=%s%s&beamLineName=%s" % (code, number, beamline)
)
]
)
[docs] def find_sessions_by_proposal_and_beamline_for_today(
self, code: str, number: str, beamline: str
) -> list[Session]:
"""Finds todays sessions by proposal code, number and beamline name."""
today = datetime.today() # noqa: DTZ002
day, month, year = today.day, today.month, today.year
return [
self.__to_session(session)
for session in self.client.get(
"sessions?proposal=%s%s&beamLineName=%s&year=%s&month=%s&day=%s"
% (code, number, beamline, year, month, day)
)
if self.__is_time_between(
datetime.fromisoformat(session.get("startDate")),
datetime.fromisoformat(session.get("endDate")),
)
]
[docs] def get_sessions_by_username(
self,
username: str = "",
beamline_name: str = "",
) -> LimsSessionManager:
"""Get the list of sessions for the authenticated user and current beamline.
PyISPyB returns only proposals accessible to the authenticated user.
For each proposal, the method fetches sessions for the current month and
picks one overlapping with the current time. If no such session exists,
a new one is created.
Args:
username: Username to fetch sessions for (left for consistency)
beamline_name: Beamline name to fetch sessions for (left for consistency)
Returns:
LimsSessionManager: A manager containing the list of sessions
"""
sessions: list[Session] = []
for proposal in self.get_proposals():
try:
session = self.find_sessions_by_proposal_and_beamline_for_today(
proposal.code, proposal.number, self.beamline_name
)[0]
except IndexError:
self.logger.info(
"No sessions planned for proposal %s%s. Creating new session.",
proposal.code,
proposal.number,
)
session = self.create_session(proposal)
sessions.append(session)
return LimsSessionManager(sessions=sessions)
[docs] def create_session(self, proposal: Proposal) -> Session:
"""Creates new session via PyISPyB REST API for the given proposal."""
start_time = datetime.today().replace(hour=0, minute=0, second=0, microsecond=0) # noqa: DTZ002
end_time = start_time + timedelta(
days=self.new_session_duration_days, hours=7, minutes=59, seconds=59
)
payload = {
"proposalId": proposal.proposal_id,
"startDate": start_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
"endDate": end_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
"beamLineName": self.beamline_name,
"comments": "Session created by the BCM",
"scheduled": False,
}
return self.__to_session(self.client.post("sessions", json=payload))
def _store_data_collection_group(self, group_data: dict) -> dict:
"""Stores data collection group in PyISPyB and returns the group id.
Args:
group_data: A dictionary containing data collection group information.
Returns:
Response from PyISPyB: {
"dataCollectionGroupId": 0,
"experimentType": "string",
"blSampleId": 0,
"sessionId": 0,
"workflowId": 0
}
"""
return self.client.post("datacollections/group", json=group_data)
def _update_data_collection(self, mx_collection: dict) -> tuple[int, int]:
"""Updates data collection in PyISPyB.
Args:
mx_collection: A dictionary containing data collection information.
Expected to contain "collection_id" key.
Returns:
Tuple containing data collection id and data collection group id.
"""
dc_id, dc_group_id = 0, 0
if "collection_id" in mx_collection:
group_id = self._store_data_collection_group(mx_collection)[
"dataCollectionGroupId"
]
try:
dc_id, dc_group_id, *_ = self.client.post(
"datacollections/datacollection",
json={**mx_collection, "group_id": group_id},
)
except PyISPyBUnsuccessfulResponse:
self.logger.exception("Error in _update_data_collection")
else:
self.logger.error(
"Error in _update_data_collection: collection-id missing, "
"the PyISPyB data-collection is not updated."
)
return dc_id, dc_group_id
[docs] def store_image(self, image_dict: dict) -> int:
"""Stores a data collection image in PyISPyB.
Args:
image_dict: A dictionary containing image information (paths and properties
of data collection). Expected to contain "dataCollectionId" key.
Returns:
The id of the stored image. 0 if failed or dataCollectionId is missing.
"""
self.logger.info("Storing image in LIMS")
image_id = 0
if "dataCollectionId" in image_dict:
try:
response = self.client.post("images/image", json=image_dict)
self.logger.debug("PyISPYB store image response: %s", response)
image_id = response["imageId"]
except PyISPyBUnsuccessfulResponse:
self.logger.exception("Exception in store_image")
else:
self.logger.info("Storing image in LIMS completed, id : %s", image_id)
else:
self.logger.error(
"Error in store_image: data_collection_id missing, "
"could not store image in LIMS"
)
return image_id
[docs] def get_samples(self, proposal_id: int) -> list[dict]:
"""Fetches samples for the given proposal id from PyISPyB."""
try:
proposal = self.__find_proposal_by_id(proposal_id)
samples = self.client.get(
"samples?proposal=%s%s&beamLineName=%s"
% (proposal.code, proposal.number, self.beamline_name),
timeout=10,
)
except PyISPyBUnsuccessfulResponse:
self.logger.exception("Error in get_samples")
return []
return [self.__to_sample(sample) for sample in samples]
[docs] def store_robot_action(self, robot_action: dict) -> int:
"""Stores robot action.
Args:
robot_action: A dictionary containing robot action information.
Expected to contain "sampleId" key.
Returns:
The id of the robot action. 0 if failed or required keys are missing.
"""
self.logger.info("Storing robot actions in LIMS")
robot_action_id = 0
try:
robot_action["blSampleId"] = robot_action.pop("sampleId")
except KeyError:
self.logger.exception(
"Error in store_robot_action: missing required keys in robot_action"
"dictionary, could not store robot action in LIMS"
)
else:
robot_action["startTimestamp"] = robot_action.pop("startTime", "")
robot_action["endTimestamp"] = robot_action.pop("endTime", "")
try:
robot_action_id = self.client.post(
"events/robot-action", json=robot_action
)
except PyISPyBUnsuccessfulResponse:
self.logger.exception("Exception in store_robot_action")
return robot_action_id
[docs] def associate_bl_sample_and_energy_scan(self, entry_dict: dict) -> int:
"""Associates bl sample and energy scan in PyISPyB.
Args:
entry_dict: A dictionary containing energy scan and bl sample information.
Returns:
The id of the association. -1 if association failed.
"""
payload = {
"energyScanId": entry_dict["energyScanId"],
"blSampleId": entry_dict["blSampleId"],
}
try:
assoc_id = self.client.patch(
"events/energyscan/associate-bl-sample", json=payload
)
except Exception:
self.logger.exception(
"Failed to associate bl sample and energy scan in PyISPyB"
)
assoc_id = -1
return assoc_id
[docs] def get_data_collection(self, data_collection_id: int) -> dict:
"""Fetches data collection details from PyISPyB.
Args:
data_collection_id: The id of the data collection to fetch.
Returns:
Dictionary containing data collection details or empty if fetching failed.
"""
try:
dc = self.client.get(f"datacollections/{data_collection_id}")
except Exception:
self.logger.exception("Failed to get data collection from PyISPyB")
return {}
if not dc:
return {}
dc["startTime"] = datetime.fromisoformat(dc["startTime"]).strftime(
"%Y-%m-%d %H:%M:%S"
)
dc["endTime"] = datetime.fromisoformat(dc["endTime"]).strftime(
"%Y-%m-%d %H:%M:%S"
)
return dc
[docs] def find_detector(
self,
manufacturer,
model,
mode,
type="", # noqa: A002
) -> dict | None:
"""Finds detector by its type, manufacturer, model and mode.
Gives the first matching detector or None if no match is found.
Args:
manufacturer: The manufacturer of the detector.
model: The model of the detector.
mode: The mode of the detector.
type: The type of the detector (optional).
Returns:
A dictionary containing detector details.
"""
try:
return self.client.get(
"detectors?manufacturer=%s&model=%s&mode=%s&type=%s"
% (
manufacturer,
model,
mode,
type if type else "",
)
)
except PyISPyBUnsuccessfulResponse:
self.logger.exception("`Exception in find_detector")
return None
[docs] def update_session(self, session: dict) -> dict:
"""Updates Beamline Setup of existing session.
Returns dictionary containing updated session details from PyISPyB.
If updating failed, returns empty dictionary.
"""
try:
session_id = session["sessionId"]
beamline_setup = session.get("BeamLineSetup") or {}
beamline_setup_id = beamline_setup.get("beamLineSetupId")
if not beamline_setup_id:
msg = "Missing beamLineSetupId"
raise KeyError(msg)
return self.client.patch(
"sessions/%s/associate-beamline-setup?beamLineSetupId=%s"
% (
session_id,
beamline_setup_id,
),
)
except (PyISPyBUnsuccessfulResponse, KeyError):
self.logger.exception("Failed to store or update session")
return {}
[docs] def get_session(self, session_id: int) -> dict:
"""Fetches session details from PyISPyB by session id."""
try:
session = self.client.get("sessions/%s" % session_id)
except PyISPyBUnsuccessfulResponse:
self.logger.exception("Failed to get session from PyISPyB")
session = {}
return session
[docs] def store_beamline_setup(self, session_id: int, bl_config: dict) -> int | None:
"""
Stores beamline setup in PyISPyB and associates it with the session.
Args:
session_id: The id of the session to associate the beamline setup with.
bl_config: A dictionary containing beamline setup information.
Returns:
The id of the stored beamline setup. None if failed."""
beamline_setup_id = None
session = self.get_session(session_id)
if session:
try:
beamline_setup_id = self.client.post(
"beamline-setups/beamline-setup", json=bl_config
)
except PyISPyBUnsuccessfulResponse:
self.logger.exception(
"Failed to store or update beamline setup in PyISPyB"
)
bl_config["beamlineSetupId"] = beamline_setup_id
session["BeamLineSetup"] = bl_config
self.update_session(session)
return beamline_setup_id
def store_data_collection(self, mx_collection, bl_config=None):
self.logger.info("Storing datacollection in PyISPyB")
return self._store_data_collection(mx_collection, bl_config)
def update_data_collection(self, mx_collection):
self.logger.info("Updating datacollection in PyISPyB")
return self._update_data_collection(mx_collection)
def finalize_data_collection(self, mx_collection):
return self.update_data_collection(mx_collection)
def _store_data_collection(self, mx_collection: dict, bl_config: dict | None):
"""Stores data collection group in PyISPyB and returns the group id.
Args:
mx_collection: A dictionary containing data collection information.
bl_config: A dictionary containing beamline setup information (optional).
Returns:
Store data collection id and detector id if found, otherwise 0.
"""
detector_id = 0
if bl_config:
bl_config["synchrotronMode"] = bl_config.get(
"synchrotronMode", mx_collection.get("synchrotronMode", "")
)
self.store_beamline_setup(mx_collection["sessionId"], bl_config)
detector = self.find_detector(
bl_config.get("detector_manufacturer", ""),
bl_config.get("detector_model", ""),
bl_config.get("detector_binning_mode", ""),
bl_config.get("detector_type", ""),
)
if detector:
detector_id = detector.get("detectorId", 0)
response = self.client.post(
"datacollections/datacollection",
json={**mx_collection, "detectorId": detector_id}
if detector_id
else mx_collection,
)
collection_id = response["dataCollectionId"]
return collection_id, detector_id
[docs] def store_energy_scan(self, energyscan: dict) -> dict[str, int]:
"""Stores energy scan in PyISPyB
Args:
energyscan: A dictionary containing energy scan information.
Returns:
Dictionary {"energyScanId": number}, where number is the id of the stored
energy scan.
"""
try:
return {
"energyScanId": self.client.post("events/energyscan", json=energyscan)
}
except PyISPyBUnsuccessfulResponse:
self.logger.exception("Failed to store energy scan in PyISPyB")
return {"energyScanId": -1}
[docs] def store_xfe_spectrum(self, xfe_spectrum: dict) -> dict[str, int]:
"""Stores XFE fluorescence spectrum in PyISPyB
Args:
xfe_spectrum: A dictionary containing XFE fluorescence spectrum information.
Returns:
Dictionary {"xfeFluorescenceSpectrumId": number}, where number is the id
of the stored XFE fluorescence spectrum.
"""
try:
return {
"xfeFluorescenceSpectrumId": self.client.post(
"events/xfe-fluorescence-spectrum", json=xfe_spectrum
)
}
except PyISPyBUnsuccessfulResponse:
self.logger.exception(
"Failed to store XFE fluorescence spectrum in PyISPyB"
)
return {"xfeFluorescenceSpectrumId": -1}
[docs] def update_bl_sample(self, bl_sample: dict) -> dict:
"""Updates beamline sample in PyISPyB.
Args:
bl_sample: A dictionary containing beamline sample information.
Expected to contain "blSampleId" key.
Returns:
Response from PyISPyB - a dictionary containing updated beamline sample
details or empty if updating failed.
"""
try:
bl_sample_id = bl_sample.get("blSampleId")
if bl_sample_id is None:
self.logger.error("Missing blSampleId")
return {}
return self.client.patch("samples/%s" % bl_sample_id, json=bl_sample)
except PyISPyBUnsuccessfulResponse:
self.logger.exception("Failed to update beamline sample in PyISPyB")
return {}