diff --git a/ami/main/management/commands/import_taxa.py b/ami/main/management/commands/import_taxa.py index 73eb1cb08..a1f9cf49b 100644 --- a/ami/main/management/commands/import_taxa.py +++ b/ami/main/management/commands/import_taxa.py @@ -166,7 +166,9 @@ class Command(BaseCommand): This is a very specific command for importing taxa from an exiting format. A more general import command with support for all taxon ranks & fields should be written. - + @TODO: Add --project parameter(s) to scope the taxa list to specific projects. + This would allow multiple projects to have taxa lists with the same name. + Usage would be: --project project-slug --project another-project-slug Example taxa.json ``` @@ -234,7 +236,8 @@ def handle(self, *args, **options): else: list_name = pathlib.Path(fname).stem - taxalist, created = TaxaList.objects.get_or_create(name=list_name) + # Uses get_or_create_for_project with project=None to create a global list + taxalist, created = TaxaList.objects.get_or_create_for_project(name=list_name, project=None) if created: self.stdout.write(self.style.SUCCESS('Successfully created taxa list "%s"' % taxalist)) diff --git a/ami/main/management/commands/update_taxa.py b/ami/main/management/commands/update_taxa.py index 4e1d75c33..92779fcbc 100644 --- a/ami/main/management/commands/update_taxa.py +++ b/ami/main/management/commands/update_taxa.py @@ -91,6 +91,10 @@ class Command(BaseCommand): ``` You can include any column that exists in the Taxon model. + + @TODO: Add --project parameter(s) to scope the taxa list to specific projects. + This would allow multiple projects to have taxa lists with the same name. + Usage would be: --project project-slug --project another-project-slug """ help = "Update existing taxa with new data from a CSV file." @@ -123,10 +127,11 @@ def handle(self, *args, **options): incoming_taxa = read_csv(fname) # Get or create taxa list if specified + # Uses get_or_create_for_project with project=None to create a global list taxalist = None if options["list"]: list_name = options["list"] - taxalist, created = TaxaList.objects.get_or_create(name=list_name) + taxalist, created = TaxaList.objects.get_or_create_for_project(name=list_name, project=None) if created: self.stdout.write(self.style.SUCCESS(f"Created new taxa list '{list_name}'")) else: diff --git a/ami/main/models.py b/ami/main/models.py index 6f67b3e3c..62c0c80ff 100644 --- a/ami/main/models.py +++ b/ami/main/models.py @@ -3595,6 +3595,44 @@ def save(self, update_calculated_fields=True, *args, **kwargs): self.update_calculated_fields(save=True) +class TaxaListQuerySet(BaseQuerySet): + def get_or_create_for_project( + self, name: str, project: "Project | None" = None, **defaults + ) -> tuple["TaxaList", bool]: + """ + Get or create a TaxaList with uniqueness scoped to project. + + - If project is None: looks for/creates a global list (no project associations) + - If project is provided: looks for/creates a list associated with that project + + Returns: + Tuple of (TaxaList, created: bool) + """ + if project is None: + # Global list: find list with this name that has no project associations + qs = self.filter(name=name).annotate(project_count=models.Count("projects")).filter(project_count=0) + else: + # Project-specific: find list with this name in this project + qs = self.filter(name=name, projects=project) + + try: + return qs.get(), False + except self.model.DoesNotExist: + taxa_list = self.create(name=name, **defaults) + if project: + taxa_list.projects.add(project) + return taxa_list, True + except self.model.MultipleObjectsReturned: + # Handle existing duplicates gracefully - return the oldest one + taxa_list = qs.order_by("created_at").first() + assert taxa_list is not None # We know there's at least one + return taxa_list, False + + +class TaxaListManager(models.Manager.from_queryset(TaxaListQuerySet)): + pass + + @final class TaxaList(BaseModel): """A checklist of taxa""" @@ -3605,6 +3643,8 @@ class TaxaList(BaseModel): taxa = models.ManyToManyField(Taxon, related_name="lists") projects = models.ManyToManyField("Project", related_name="taxa_lists") + objects: TaxaListManager = TaxaListManager() + class Meta: ordering = ["-created_at"] verbose_name_plural = "Taxa Lists" diff --git a/ami/ml/models/pipeline.py b/ami/ml/models/pipeline.py index f76822e3a..6ef491730 100644 --- a/ami/ml/models/pipeline.py +++ b/ami/ml/models/pipeline.py @@ -555,8 +555,9 @@ def get_or_create_taxon_for_classification( :return: The Taxon object """ - taxa_list, created = TaxaList.objects.get_or_create( + taxa_list, created = TaxaList.objects.get_or_create_for_project( name=f"Taxa returned by {algorithm.name}", + project=None, # Algorithm taxa lists are global ) if created: logger.info(f"Created new taxa list {taxa_list}")