from typing import List, Optional from datetime import datetime from pydantic import BaseModel, EmailStr, constr, Field, field_validator from datetime import date import logging import re logging.basicConfig(level=logging.DEBUG) logger = logging.getLogger(__name__) class loginToken(BaseModel): access_token: str token_type: str class loginData(BaseModel): username: str pgroups: List[str] class DewarTypeBase(BaseModel): dewar_type: str class DewarTypeCreate(DewarTypeBase): pass class DewarType(DewarTypeBase): id: int class Config: from_attributes = True class DewarSerialNumberBase(BaseModel): serial_number: str dewar_type_id: int class DewarSerialNumberCreate(DewarSerialNumberBase): pass class DewarSerialNumber(DewarSerialNumberBase): id: int dewar_type: DewarType class Config: from_attributes = True class DataCollectionParameters(BaseModel): directory: Optional[str] = None oscillation: Optional[float] = None # Only accept positive float exposure: Optional[float] = None # Only accept positive floats between 0 and 1 totalrange: Optional[int] = None # Only accept positive integers between 0 and 360 transmission: Optional[ int ] = None # Only accept positive integers between 0 and 100 targetresolution: Optional[float] = None # Only accept positive float aperture: Optional[str] = None # Optional string field datacollectiontype: Optional[ str ] = None # Only accept "standard", other types might be added later processingpipeline: Optional[ str ] = "" # Only accept "gopy", "autoproc", "xia2dials" spacegroupnumber: Optional[ int ] = None # Only accept positive integers between 1 and 230 cellparameters: Optional[ str ] = None # Must be a set of six positive floats or integers rescutkey: Optional[str] = None # Only accept "is" or "cchalf" rescutvalue: Optional[ float ] = None # Must be a positive float if rescutkey is provided userresolution: Optional[float] = None pdbid: Optional[ str ] = "" # Accepts either the format of the protein data bank code or {provided} autoprocfull: Optional[bool] = None procfull: Optional[bool] = None adpenabled: Optional[bool] = None noano: Optional[bool] = None ffcscampaign: Optional[bool] = None trustedhigh: Optional[float] = None # Should be a float between 0 and 2.0 autoprocextraparams: Optional[str] = None # Optional string field chiphiangles: Optional[float] = None # Optional float field between 0 and 30 dose: Optional[float] = None # Optional float field def to_dict(self): """Convert the model instance to a dictionary.""" return self.dict( exclude_unset=True ) # Use this built-in method for serialization class Config: from_attributes = True @field_validator("directory", mode="after") @classmethod def directory_characters(cls, v): logger.debug(f"Validating 'directory' field with initial value: {repr(v)}") # Default directory value if empty if not v: # Handles None or empty cases default_value = "{sgPuck}/{sgPosition}" logger.warning( f"'directory' field is empty or None. Assigning default value: " f"{default_value}" ) return default_value # Strip trailing slashes and store original value for comparison v = str(v).strip("/") # Ensure it's a string and no trailing slashes original_value = v # Replace spaces with underscores v = v.replace(" ", "_") logger.debug(f"Corrected 'directory', spaces replaced: {repr(v)}") # Validate directory pattern with macros and allowed characters valid_macros = [ "{date}", "{prefix}", "{sgPuck}", "{sgPosition}", "{beamline}", "{sgPrefix}", "{sgPriority}", "{protein}", "{method}", ] valid_macro_pattern = re.compile( "|".join(re.escape(macro) for macro in valid_macros) ) # Check if the value contains valid macros allowed_chars_pattern = "[a-z0-9_.+-/]" v_without_macros = valid_macro_pattern.sub("macro", v) allowed_path_pattern = re.compile( f"^(({allowed_chars_pattern}+|macro)*/*)*$", re.IGNORECASE ) if not allowed_path_pattern.match(v_without_macros): raise ValueError( f"'{v}' is not valid. Value must be a valid path or macro." ) # Log and return corrected value if v != original_value: logger.info(f"Directory was corrected from '{original_value}' to '{v}'") return v @field_validator("aperture", mode="before") @classmethod def aperture_selection(cls, v): if v is not None: try: v = int(float(v)) if v not in {1, 2, 3}: raise ValueError(f" '{v}' is not valid. Value must be 1, 2, or 3.") except (ValueError, TypeError) as e: raise ValueError( f" '{v}' is not valid. Value must be 1, 2, or 3." ) from e return v @field_validator("oscillation", mode="before") @classmethod def positive_float_validator(cls, v): if v is None: return None try: v = float(v) if v <= 0: raise ValueError(f"'{v}' is not valid. Value must be a positive float.") except (ValueError, TypeError) as e: raise ValueError( f"'{v}' is not valid. Value must be a positive float." ) from e return v @field_validator("exposure", mode="before") @classmethod def exposure_in_range(cls, v): if v is not None: try: v = float(v) if not (0 <= v <= 1): raise ValueError( f" '{v}' is not valid. Value must be a float between 0 and 1." ) except (ValueError, TypeError) as e: raise ValueError( f" '{v}' is not valid. Value must be a float between 0 and 1." ) from e return v @field_validator("totalrange", mode="before") @classmethod def totalrange_in_range(cls, v): if v is not None: try: v = int(v) if not (0 <= v <= 360): raise ValueError( f" '{v}' is not valid." f"Value must be an integer between 0 and 360." ) except (ValueError, TypeError) as e: raise ValueError( f" '{v}' is not valid." f"Value must be an integer between 0 and 360." ) from e return v @field_validator("transmission", mode="before") @classmethod def transmission_fraction(cls, v): if v is not None: try: v = int(v) if not (0 <= v <= 100): raise ValueError( f" '{v}' is not valid." f"Value must be an integer between 0 and 100." ) except (ValueError, TypeError) as e: raise ValueError( f" '{v}' is not valid." f"Value must be an integer between 0 and 100." ) from e return v @field_validator("datacollectiontype", mode="before") @classmethod def datacollectiontype_allowed(cls, v): allowed = {"standard"} # Other types of data collection might be added later if v and v.lower() not in allowed: raise ValueError(f" '{v}' is not valid." f"Value must be one of {allowed}.") return v @field_validator("processingpipeline", mode="before") @classmethod def processingpipeline_allowed(cls, v): allowed = {"gopy", "autoproc", "xia2dials"} if v and v.lower() not in allowed: raise ValueError(f" '{v}' is not valid." f"Value must be one of {allowed}.") return v @field_validator("spacegroupnumber", mode="before") @classmethod def spacegroupnumber_allowed(cls, v): if v is not None: try: v = int(v) if not (1 <= v <= 230): raise ValueError( f" '{v}' is not valid." f"Value must be an integer between 1 and 230." ) except (ValueError, TypeError) as e: raise ValueError( f" '{v}' is not valid." f"Value must be an integer between 1 and 230." ) from e return v @field_validator("cellparameters", mode="before") @classmethod def cellparameters_format(cls, v): if v: values = [float(i) for i in v.split(",")] if len(values) != 6 or any(val <= 0 for val in values): raise ValueError( f" '{v}' is not valid." f"Value must be a set of six positive floats or integers." ) return v # @field_validator("rescutkey", "rescutvalue", mode="before") # @classmethod # def rescutkey_value_pair(cls, values): # rescutkey = values.get("rescutkey") # rescutvalue = values.get("rescutvalue") # if rescutkey and rescutvalue: # if rescutkey not in {"is", "cchalf"}: # raise ValueError("Rescutkey must be either 'is' or 'cchalf'") # if not isinstance(rescutvalue, float) or rescutvalue <= 0: # raise ValueError( # "Rescutvalue must be a positive float if rescutkey is provided" # ) # return values @field_validator("trustedhigh", mode="before") @classmethod def trustedhigh_allowed(cls, v): if v is not None: try: v = float(v) if not (0 <= v <= 2.0): raise ValueError( f" '{v}' is not valid." f"Value must be a float between 0 and 2.0." ) except (ValueError, TypeError) as e: raise ValueError( f" '{v}' is not valid." f"Value must be a float between 0 and 2.0." ) from e return v @field_validator("chiphiangles", mode="before") @classmethod def chiphiangles_allowed(cls, v): if v is not None: try: v = float(v) if not (0 <= v <= 30): raise ValueError( f" '{v}' is not valid." f"Value must be a float between 0 and 30." ) except (ValueError, TypeError) as e: raise ValueError( f" '{v}' is not valid. Value must be a float between 0 and 30." ) from e return v @field_validator("dose", mode="before") @classmethod def dose_positive(cls, v): if v is not None: try: v = float(v) if v <= 0: raise ValueError( f" '{v}' is not valid. Value must be a positive float." ) except (ValueError, TypeError) as e: raise ValueError( f" '{v}' is not valid. Value must be a positive float." ) from e return v class SampleEventCreate(BaseModel): event_type: str class SampleEventResponse(BaseModel): id: int sample_id: int event_type: str timestamp: datetime class Config: from_attributes = True class Results(BaseModel): # Define attributes for Results here pass class ContactCreate(BaseModel): pgroups: str firstname: str lastname: str phone_number: str email: EmailStr class Contact(ContactCreate): id: int status: str = "active" class Config: from_attributes = True class ContactUpdate(BaseModel): pgroups: str firstname: Optional[str] = None lastname: Optional[str] = None phone_number: Optional[str] = None email: Optional[EmailStr] = None class AddressCreate(BaseModel): pgroups: str house_number: Optional[str] = None street: str city: str state: Optional[str] = None zipcode: str country: str class Address(AddressCreate): id: int status: str = "active" class Config: from_attributes = True class AddressUpdate(BaseModel): pgroups: str house_number: Optional[str] = None street: Optional[str] = None city: Optional[str] = None state: Optional[str] = None zipcode: Optional[str] = None country: Optional[str] = None class Sample(BaseModel): id: int sample_name: str position: int # Position within the puck puck_id: int crystalname: Optional[str] = Field(None) proteinname: Optional[str] = None positioninpuck: Optional[int] = Field(None) priority: Optional[int] = None comments: Optional[str] = None data_collection_parameters: Optional[DataCollectionParameters] events: List[SampleEventResponse] mount_count: Optional[int] = None unmount_count: Optional[int] = None # results: Optional[Results] = None class Config: from_attributes = True class SampleCreate(BaseModel): sample_name: str = Field(..., alias="crystalname") proteinname: Optional[str] = None position: int = Field(..., alias="positioninpuck") data_collection_parameters: Optional[DataCollectionParameters] = Field(default=None) priority: Optional[int] = None comments: Optional[str] = None results: Optional[Results] = None events: Optional[List[str]] = None class Config: populate_by_name = True class PuckEvent(BaseModel): id: int puck_id: int tell_position: Optional[str] = None event_type: str timestamp: datetime class Config: from_attributes = True class PuckBase(BaseModel): puck_name: str puck_type: str puck_location_in_dewar: int class PuckCreate(BaseModel): puck_name: str puck_type: str puck_location_in_dewar: int samples: List[SampleCreate] = [] class PuckUpdate(BaseModel): puck_name: Optional[str] = None puck_type: Optional[str] = None puck_location_in_dewar: Optional[int] = None dewar_id: Optional[int] = None class Puck(BaseModel): id: int puck_name: str puck_type: str puck_location_in_dewar: int dewar_id: int events: List[PuckEvent] = [] samples: List[Sample] = [] class Config: from_attributes = True class DewarBase(BaseModel): dewar_name: str dewar_type_id: Optional[int] = None dewar_serial_number_id: Optional[int] = None unique_id: Optional[str] = None tracking_number: str number_of_pucks: Optional[int] = None number_of_samples: Optional[int] = None status: str ready_date: Optional[date] shipping_date: Optional[date] arrival_date: Optional[date] returning_date: Optional[date] contact_id: Optional[int] return_address_id: Optional[int] pucks: List[PuckCreate] = [] class Config: from_attributes = True class DewarCreate(DewarBase): pass class Dewar(DewarBase): id: int shipment_id: Optional[int] contact: Optional[Contact] return_address: Optional[Address] pucks: List[Puck] = [] # List of pucks within this dewar class Config: from_attributes = True class DewarUpdate(BaseModel): dewar_name: Optional[str] = None dewar_type_id: Optional[int] = None dewar_serial_number_id: Optional[int] = None unique_id: Optional[str] = None tracking_number: Optional[str] = None status: Optional[str] = None ready_date: Optional[date] = None shipping_date: Optional[date] = None arrival_date: Optional[date] = None returning_date: Optional[date] = None contact_id: Optional[int] = None address_id: Optional[int] = None class DewarSchema(BaseModel): id: int dewar_name: str tracking_number: str status: str contact_id: int return_address_id: int class Config: from_attributes = True class Proposal(BaseModel): id: int number: str class Config: from_attributes = True class Shipment(BaseModel): id: int pgroups: str shipment_name: str shipment_date: date shipment_status: str comments: Optional[str] contact: Optional[Contact] return_address: Optional[Address] proposal: Optional[Proposal] dewars: List[Dewar] = [] class Config: from_attributes = True class ShipmentCreate(BaseModel): pgroups: str shipment_name: str shipment_date: date shipment_status: str comments: Optional[constr(max_length=200)] contact_id: int return_address_id: int proposal_id: int dewars: List[DewarCreate] = [] class Config: from_attributes = True class UpdateShipmentComments(BaseModel): pgroups: str comments: str class LogisticsEventCreate(BaseModel): dewar_qr_code: str location_qr_code: str transaction_type: str class SlotSchema(BaseModel): id: int qr_code: str label: str qr_base: Optional[str] occupied: bool needs_refill: bool dewar_unique_id: Optional[str] dewar_name: Optional[str] time_until_refill: Optional[int] at_beamline: Optional[bool] retrievedTimestamp: Optional[str] beamlineLocation: Optional[str] shipment_name: Optional[str] contact: Optional[str] local_contact: Optional[str] class Config: from_attributes = True class SampleUpdate(BaseModel): sample_name: Optional[str] = None proteinname: Optional[str] = None priority: Optional[int] = None position: Optional[int] = None comments: Optional[str] = None data_collection_parameters: Optional[DataCollectionParameters] = None class Config: from_attributes = True class SetTellPosition(BaseModel): puck_name: str segment: Optional[str] = Field( None, pattern="^[A-F]$", # Valid segments are A, B, C, D, E, F description="Segment must be one of A, B, C, D, E, or F." "Can be null for no tell_position.", ) puck_in_segment: Optional[int] = Field( None, ge=1, le=5, description="Puck in segment must be between 1 and 5." "Can be null for no tell_position.", ) @property def tell_position(self) -> Optional[str]: """ Combines `segment` and `puck_in_segment` to generate the `tell_position`. If either value is `None`, returns `None` to indicate no `tell_position`. """ if self.segment and self.puck_in_segment: return f"{self.segment}{self.puck_in_segment}" return None class PuckWithTellPosition(BaseModel): id: int puck_name: str puck_type: str puck_location_in_dewar: Optional[int] dewar_id: Optional[ int ] # was changed to optional but probably needs to be not optional dewar_name: Optional[ str ] # was changed to optional but probably needs to be not optional user: str = "e16371" samples: Optional[List[Sample]] = None tell_position: Optional[str] class Config: from_attributes = True