Creation of new scripts.

src/submissions/backend/db/models/kits.py

@@ -2,12 +2,12 @@
 All kit and reagent related models
 """
 from __future__ import annotations
-import datetime, json, zipfile, yaml, logging, re
+import json, zipfile, yaml, logging, re
 from pprint import pformat
 from sqlalchemy import Column, String, TIMESTAMP, JSON, INTEGER, ForeignKey, Interval, Table, FLOAT, BLOB
 from sqlalchemy.orm import relationship, validates, Query
 from sqlalchemy.ext.associationproxy import association_proxy
-from datetime import date, datetime
+from datetime import date, datetime, timedelta
 from tools import check_authorization, setup_lookup, Report, Result, check_regex_match, yaml_regex_creator, timezone
 from typing import List, Literal, Generator, Any
 from pandas import ExcelFile
@@ -259,6 +259,79 @@ class KitType(BaseClass):
             base_dict['equipment_roles'].append(v)
         return base_dict
 
+    @classmethod
+    def import_from_yml(cls, submission_type:str|SubmissionType, filepath: Path | str | None = None, import_dict: dict | None = None) -> KitType:
+        if isinstance(submission_type, str):
+            submission_type = SubmissionType.query(name=submission_type)
+        if filepath:
+            yaml.add_constructor("!regex", yaml_regex_creator)
+            if isinstance(filepath, str):
+                filepath = Path(filepath)
+            if not filepath.exists():
+                logging.critical(f"Given file could not be found.")
+                return None
+            with open(filepath, "r") as f:
+                if filepath.suffix == ".json":
+                    import_dict = json.load(fp=f)
+                elif filepath.suffix == ".yml":
+                    import_dict = yaml.load(stream=f, Loader=yaml.Loader)
+                else:
+                    raise Exception(f"Filetype {filepath.suffix} not supported.")
+        new_kit = KitType.query(name=import_dict['kit_type']['name'])
+        if not new_kit:
+            new_kit = KitType(name=import_dict['kit_type']['name'])
+        for role in import_dict['kit_type']['reagent_roles']:
+            new_role = ReagentRole.query(name=role['role'])
+            if new_role:
+                check = input(f"Found existing role: {new_role.name}. Use this? [Y/n]: ")
+                if check.lower() == "n":
+                    new_role = None
+                else:
+                    pass
+            if not new_role:
+                eol = timedelta(role['extension_of_life'])
+                new_role = ReagentRole(name=role['role'], eol_ext=eol)
+            uses = dict(expiry=role['expiry'], lot=role['lot'], name=role['name'], sheet=role['sheet'])
+            ktrr_assoc = KitTypeReagentRoleAssociation(kit_type=new_kit, reagent_role=new_role, uses=uses)
+            ktrr_assoc.submission_type = submission_type
+            ktrr_assoc.required = role['required']
+        ktst_assoc = SubmissionTypeKitTypeAssociation(
+            kit_type=new_kit,
+            submission_type=submission_type,
+            mutable_cost_sample=import_dict['mutable_cost_sample'],
+            mutable_cost_column=import_dict['mutable_cost_column'],
+            constant_cost=import_dict['constant_cost']
+        )
+        for role in import_dict['kit_type']['equipment_roles']:
+            new_role = EquipmentRole.query(name=role['role'])
+            if new_role:
+                check = input(f"Found existing role: {new_role.name}. Use this? [Y/n]: ")
+                if check.lower() == "n":
+                    new_role = None
+                else:
+                    pass
+            if not new_role:
+                new_role = EquipmentRole(name=role['role'])
+                for equipment in Equipment.assign_equipment(equipment_role=new_role):
+                    new_role.instances.append(equipment)
+            ster_assoc = SubmissionTypeEquipmentRoleAssociation(submission_type=submission_type,
+                                                                equipment_role=new_role)
+            try:
+                uses = dict(name=role['name'], process=role['process'], sheet=role['sheet'],
+                            static=role['static'])
+            except KeyError:
+                uses = None
+            ster_assoc.uses = uses
+            for process in role['processes']:
+                new_process = Process.query(name=process)
+                if not new_process:
+                    new_process = Process(name=process)
+                new_process.submission_types.append(submission_type)
+                new_process.kit_types.append(new_kit)
+                new_process.equipment_roles.append(new_role)
+        return new_kit
+
+
 
 class ReagentRole(BaseClass):
     """
@@ -903,58 +976,7 @@ class SubmissionType(BaseClass):
             submission_type.sample_map = import_dict['samples']
             submission_type.defaults = import_dict['defaults']
             for kit in import_dict['kits']:
-                new_kit = KitType.query(name=kit['kit_type']['name'])
-                if not new_kit:
-                    new_kit = KitType(name=kit['kit_type']['name'])
-                for role in kit['kit_type']['reagent roles']:
-                    new_role = ReagentRole.query(name=role['role'])
-                    if new_role:
-                        check = input(f"Found existing role: {new_role.name}. Use this? [Y/n]: ")
-                        if check.lower() == "n":
-                            new_role = None
-                        else:
-                            pass
-                    if not new_role:
-                        eol = datetime.timedelta(role['extension_of_life'])
-                        new_role = ReagentRole(name=role['role'], eol_ext=eol)
-                    uses = dict(expiry=role['expiry'], lot=role['lot'], name=role['name'], sheet=role['sheet'])
-                    ktrr_assoc = KitTypeReagentRoleAssociation(kit_type=new_kit, reagent_role=new_role, uses=uses)
-                    ktrr_assoc.submission_type = submission_type
-                    ktrr_assoc.required = role['required']
-                ktst_assoc = SubmissionTypeKitTypeAssociation(
-                    kit_type=new_kit,
-                    submission_type=submission_type,
-                    mutable_cost_sample=kit['mutable_cost_sample'],
-                    mutable_cost_column=kit['mutable_cost_column'],
-                    constant_cost=kit['constant_cost']
-                )
-                for role in kit['kit_type']['equipment roles']:
-                    new_role = EquipmentRole.query(name=role['role'])
-                    if new_role:
-                        check = input(f"Found existing role: {new_role.name}. Use this? [Y/n]: ")
-                        if check.lower() == "n":
-                            new_role = None
-                        else:
-                            pass
-                    if not new_role:
-                        new_role = EquipmentRole(name=role['role'])
-                        for equipment in Equipment.assign_equipment(equipment_role=new_role):
-                            new_role.instances.append(equipment)
-                    ster_assoc = SubmissionTypeEquipmentRoleAssociation(submission_type=submission_type,
-                                                                        equipment_role=new_role)
-                    try:
-                        uses = dict(name=role['name'], process=role['process'], sheet=role['sheet'],
-                                    static=role['static'])
-                    except KeyError:
-                        uses = None
-                    ster_assoc.uses = uses
-                    for process in role['processes']:
-                        new_process = Process.query(name=process)
-                        if not new_process:
-                            new_process = Process(name=process)
-                        new_process.submission_types.append(submission_type)
-                        new_process.kit_types.append(new_kit)
-                        new_process.equipment_roles.append(new_role)
+                new_kit = KitType.import_from_yml(submission_type=submission_type, import_dict=kit)
         if 'orgs' in import_dict.keys():
             logger.info("Found Organizations to be imported.")
             Organization.import_from_yml(filepath=filepath)
@@ -1321,7 +1343,8 @@ class Equipment(BaseClass, LogMixin):
         else:
             return {k: v for k, v in self.__dict__.items()}
 
-    def get_processes(self, submission_type: str | SubmissionType | None = None, extraction_kit: str | KitType | None = None) -> List[str]:
+    def get_processes(self, submission_type: str | SubmissionType | None = None,
+                      extraction_kit: str | KitType | None = None) -> List[str]:
         """
         Get all processes associated with this Equipment for a given SubmissionType
 
@@ -1399,7 +1422,7 @@ class Equipment(BaseClass, LogMixin):
         from backend.validators.pydant import PydEquipment
         processes = self.get_processes(submission_type=submission_type, extraction_kit=extraction_kit)
         return PydEquipment(processes=processes, role=role,
-            **self.to_dict(processes=False))
+                            **self.to_dict(processes=False))
 
     @classmethod
     def get_regex(cls) -> re.Pattern:
@@ -1547,9 +1570,9 @@ class EquipmentRole(BaseClass):
             extraction_kit = KitType.query(name=extraction_kit)
         for process in self.processes:
             if submission_type and submission_type not in process.submission_types:
-                    continue
+                continue
             if extraction_kit and extraction_kit not in process.kit_types:
-                    continue
+                continue
             yield process.name
 
     def to_export_dict(self, submission_type: SubmissionType, kit_type: KitType):
@@ -1597,7 +1620,6 @@ class SubmissionEquipmentAssociation(BaseClass):
         Returns:
             dict: This SubmissionEquipmentAssociation as a dictionary
         """
-        # TODO: Currently this will only fetch a single process, even if multiple are selectable.
         try:
             process = self.process.name
         except AttributeError:
@@ -1606,7 +1628,13 @@ class SubmissionEquipmentAssociation(BaseClass):
                       processes=[process], role=self.role, nickname=self.equipment.nickname)
         return output
 
-    def to_pydantic(self):
+    def to_pydantic(self) -> "PydEquipment":
+        """
+        Returns a pydantic model based on this object.
+
+        Returns:
+            PydEquipment: pydantic equipment model
+        """
         from backend.validators import PydEquipment
         return PydEquipment(**self.to_sub_dict())
 

src/submissions/backend/db/models/submissions.py

@@ -2,6 +2,8 @@
 Models for the main submission and sample types.
 """
 from __future__ import annotations
+
+from collections import OrderedDict
 from copy import deepcopy
 from getpass import getuser
 import logging, uuid, tempfile, re, base64, numpy as np, pandas as pd, types, sys
@@ -559,7 +561,7 @@ class BasicSubmission(BaseClass, LogMixin):
             except AttributeError as e:
                 logger.error(f"Could not set {self} attribute {key} to {value} due to \n{e}")
 
-    def update_subsampassoc(self, sample: BasicSample, input_dict: dict):
+    def update_subsampassoc(self, sample: BasicSample, input_dict: dict) -> SubmissionSampleAssociation:
         """
         Update a joined submission sample association.
 
@@ -568,7 +570,7 @@ class BasicSubmission(BaseClass, LogMixin):
             input_dict (dict): values to be updated
 
         Returns:
-            Result: _description_
+            SubmissionSampleAssociation: Updated association
         """
         try:
             assoc = next(item for item in self.submission_sample_associations if item.sample == sample)
@@ -583,14 +585,14 @@ class BasicSubmission(BaseClass, LogMixin):
                 # NOTE: for some reason I don't think assoc.__setattr__(k, v) works here.
             except AttributeError:
                 logger.error(f"Can't set {k} to {v}")
-        result = assoc.save()
-        return result
+        return assoc
 
     def update_reagentassoc(self, reagent: Reagent, role: str):
         from backend.db import SubmissionReagentAssociation
         # NOTE: get the first reagent assoc that fills the given role.
         try:
-            assoc = next(item for item in self.submission_reagent_associations if item.reagent and role in [role.name for role in item.reagent.role])
+            assoc = next(item for item in self.submission_reagent_associations if
+                         item.reagent and role in [role.name for role in item.reagent.role])
             assoc.reagent = reagent
         except StopIteration as e:
             logger.error(f"Association for {role} not found, creating new association.")
@@ -611,7 +613,8 @@ class BasicSubmission(BaseClass, LogMixin):
             missing = value in ['', 'None', None]
             match key:
                 case "reagents":
-                    field_value = [item.to_pydantic(extraction_kit=self.extraction_kit) for item in self.submission_reagent_associations]
+                    field_value = [item.to_pydantic(extraction_kit=self.extraction_kit) for item in
+                                   self.submission_reagent_associations]
                 case "samples":
                     field_value = [item.to_pydantic() for item in self.submission_sample_associations]
                 case "equipment":
@@ -643,7 +646,8 @@ class BasicSubmission(BaseClass, LogMixin):
                         continue
             new_dict[key] = field_value
         new_dict['filepath'] = Path(tempfile.TemporaryFile().name)
-        return PydSubmission(**new_dict)
+        dicto.update(new_dict)
+        return PydSubmission(**dicto)
 
     def save(self, original: bool = True):
         """
@@ -1006,7 +1010,7 @@ class BasicSubmission(BaseClass, LogMixin):
     @setup_lookup
     def query(cls,
               submission_type: str | SubmissionType | None = None,
-              submission_type_name: str|None = None,
+              submission_type_name: str | None = None,
               id: int | str | None = None,
               rsl_plate_num: str | None = None,
               start_date: date | str | int | None = None,
@@ -1287,7 +1291,7 @@ class BasicSubmission(BaseClass, LogMixin):
         writer = pyd.to_writer()
         writer.xl.save(filename=fname.with_suffix(".xlsx"))
 
-    def get_turnaround_time(self) -> Tuple[int|None, bool|None]:
+    def get_turnaround_time(self) -> Tuple[int | None, bool | None]:
         try:
             completed = self.completed_date.date()
         except AttributeError:
@@ -1295,7 +1299,8 @@ class BasicSubmission(BaseClass, LogMixin):
         return self.calculate_turnaround(start_date=self.submitted_date.date(), end_date=completed)
 
     @classmethod
-    def calculate_turnaround(cls, start_date:date|None=None, end_date:date|None=None) -> Tuple[int|None, bool|None]:
+    def calculate_turnaround(cls, start_date: date | None = None, end_date: date | None = None) -> Tuple[
+        int | None, bool | None]:
         if 'pytest' not in sys.modules:
             from tools import ctx
         else:
@@ -1499,7 +1504,7 @@ class Wastewater(BasicSubmission):
         output = []
         for sample in samples:
             # NOTE: remove '-{target}' from controls
-            sample['sample'] = re.sub('-N\\d$', '', sample['sample'])
+            sample['sample'] = re.sub('-N\\d*$', '', sample['sample'])
             # NOTE: if sample is already in output skip
             if sample['sample'] in [item['sample'] for item in output]:
                 logger.warning(f"Already have {sample['sample']}")
@@ -1509,7 +1514,7 @@ class Wastewater(BasicSubmission):
             # NOTE: Set assessment
             sample[f"{sample['target'].lower()}_status"] = sample['assessment']
             # NOTE: Get sample having other target
-            other_targets = [s for s in samples if re.sub('-N\\d$', '', s['sample']) == sample['sample']]
+            other_targets = [s for s in samples if re.sub('-N\\d*$', '', s['sample']) == sample['sample']]
             for s in other_targets:
                 sample[f"ct_{s['target'].lower()}"] = s['ct'] if isinstance(s['ct'], float) else 0.0
                 sample[f"{s['target'].lower()}_status"] = s['assessment']
@@ -1613,7 +1618,9 @@ class Wastewater(BasicSubmission):
                 sample_dict = next(item for item in pcr_samples if item['sample'] == sample.rsl_number)
             except StopIteration:
                 continue
-            self.update_subsampassoc(sample=sample, input_dict=sample_dict)
+            assoc = self.update_subsampassoc(sample=sample, input_dict=sample_dict)
+            result = assoc.save()
+            report.add_result(result)
         controltype = ControlType.query(name="PCR Control")
         submitted_date = datetime.strptime(" ".join(parser.pcr['run_start_date/time'].split(" ")[:-1]),
                                            "%Y-%m-%d %I:%M:%S %p")
@@ -1623,6 +1630,27 @@ class Wastewater(BasicSubmission):
             new_control.controltype = controltype
             new_control.submission = self
             new_control.save()
+        return report
+
+    def update_subsampassoc(self, sample: BasicSample, input_dict: dict):
+        """
+        Updates a joined submission sample association by assigning ct values to n1 or n2 based on alphabetical sorting.
+
+        Args:
+            sample (BasicSample): Associated sample.
+            input_dict (dict): values to be updated
+
+        Returns:
+            SubmissionSampleAssociation: Updated association
+        """
+        assoc = super().update_subsampassoc(sample=sample, input_dict=input_dict)
+        targets = {k: input_dict[k] for k in sorted(input_dict.keys()) if k.startswith("ct_")}
+        assert 0 < len(targets) <= 2
+        for i, v in enumerate(targets.values(), start=1):
+            update_key = f"ct_n{i}"
+            if getattr(assoc, update_key) is None:
+                setattr(assoc, update_key, v)
+        return assoc
 
 
 class WastewaterArtic(BasicSubmission):
@@ -1661,7 +1689,7 @@ class WastewaterArtic(BasicSubmission):
         else:
             output['artic_technician'] = self.artic_technician
         output['gel_info'] = self.gel_info
-        output['gel_image_path'] = self.gel_image
+        output['gel_image'] = self.gel_image
         output['dna_core_submission_number'] = self.dna_core_submission_number
         output['source_plates'] = self.source_plates
         output['artic_date'] = self.artic_date or self.submitted_date
@@ -1988,7 +2016,6 @@ class WastewaterArtic(BasicSubmission):
             egel_section = custom_fields['egel_info']
             # NOTE: print json field gel results to Egel results
             worksheet = input_excel[egel_section['sheet']]
-            # TODO: Move all this into a seperate function?
             start_row = egel_section['start_row'] - 1
             start_column = egel_section['start_column'] - 3
             for row, ki in enumerate(info['gel_info']['value'], start=1):
@@ -2003,10 +2030,10 @@ class WastewaterArtic(BasicSubmission):
                         logger.error(f"Failed {kj['name']} with value {kj['value']} to row {row}, column {column}")
         else:
             logger.warning("No gel info found.")
-        if check_key_or_attr(key='gel_image_path', interest=info, check_none=True):
+        if check_key_or_attr(key='gel_image', interest=info, check_none=True):
             worksheet = input_excel[egel_section['sheet']]
             with ZipFile(cls.__directory_path__.joinpath("submission_imgs.zip")) as zipped:
-                z = zipped.extract(info['gel_image_path']['value'], Path(TemporaryDirectory().name))
+                z = zipped.extract(info['gel_image']['value'], Path(TemporaryDirectory().name))
                 img = OpenpyxlImage(z)
                 img.height = 400  # insert image height in pixels as float or int (e.g. 305.5)
                 img.width = 600
@@ -2041,9 +2068,9 @@ class WastewaterArtic(BasicSubmission):
             headers = [item['name'] for item in base_dict['gel_info'][0]['values']]
             base_dict['headers'] = [''] * (4 - len(headers))
             base_dict['headers'] += headers
-        if check_key_or_attr(key='gel_image_path', interest=base_dict, check_none=True):
+        if check_key_or_attr(key='gel_image', interest=base_dict, check_none=True):
             with ZipFile(cls.__directory_path__.joinpath("submission_imgs.zip")) as zipped:
-                base_dict['gel_image'] = base64.b64encode(zipped.read(base_dict['gel_image_path'])).decode('utf-8')
+                base_dict['gel_image_actual'] = base64.b64encode(zipped.read(base_dict['gel_image'])).decode('utf-8')
         return base_dict, template
 
     def custom_context_events(self) -> dict: