Assign sequences to segment in command

DDMAL · Jun 3, 2024 · 781387e · 781387e
1 parent 490ce55
commit 781387e
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 3 deletions.
diff --git a/django/cantusdb_project/main_app/management/commands/assign_chants_to_segments.py b/django/cantusdb_project/main_app/management/commands/assign_chants_to_segments.py
@@ -7,7 +7,7 @@
 """
 
 from django.core.management.base import BaseCommand
-from main_app.models import Source, Chant, Segment
+from main_app.models import Source, Chant, Segment, Sequence
 
 
 class Command(BaseCommand):
@@ -18,4 +18,27 @@ def handle(self, *args, **options):
         for source in sources:
             segment = Segment.objects.get(id=source.segment_id)
             chants = Chant.objects.filter(source=source)
-            chants.update(segment=segment)
+            sequences = Sequence.objects.filter(source=source)
+            chants_count = chants.count()
+            sequences_count = sequences.count()
+            if chants_count != 0 and sequences_count != 0:
+                self.stdout.write(
+                    self.style.ERROR(
+                        f"Source {source.id} has {chants_count} chants and {sequences_count} sequences."
+                    )
+                )
+                continue
+            if chants_count > 0:
+                chants.update(segment=segment)
+                self.stdout.write(
+                    self.style.SUCCESS(
+                        f"Assigned {chants_count} chants in source {source.id} to segment {segment.id}."
+                    )
+                )
+            else:
+                sequences.update(segment=segment)
+                self.stdout.write(
+                    self.style.SUCCESS(
+                        f"Assigned {sequences_count} sequences in source {source.id} to segment {segment.id}."
+                    )
+                )
diff --git a/django/cantusdb_project/main_app/tests/test_assign_chants_to_segments.py b/django/cantusdb_project/main_app/tests/test_assign_chants_to_segments.py
@@ -1,7 +1,7 @@
 from django.test import TestCase
 from django.core.management import call_command
 
-from main_app.models import Chant
+from main_app.models import Chant, Sequence
 
 from main_app.tests.make_fakes import make_fake_source, make_fake_segment
 
@@ -12,17 +12,26 @@ def test_assign_chants_to_segments(self):
         segment_2 = make_fake_segment()
         source_1 = make_fake_source(segment=segment_1)
         source_2 = make_fake_source(segment=segment_2)
+        sequence_source = make_fake_source(segment=segment_2)
         for _ in range(5):
             Chant.objects.create(source=source_1)
         for _ in range(3):
             Chant.objects.create(source=source_2)
+        for _ in range(4):
+            Sequence.objects.create(source=sequence_source)
         all_chants = Chant.objects.all()
         for chant in all_chants:
             self.assertIsNone(chant.segment_id)
+        all_sequences = Sequence.objects.all()
+        for sequence in all_sequences:
+            self.assertIsNone(sequence.segment_id)
         call_command("assign_chants_to_segments")
         source_1_chants = Chant.objects.filter(source=source_1)
         source_2_chants = Chant.objects.filter(source=source_2)
+        sequence_source_sequences = Sequence.objects.filter()
         for chant in source_1_chants:
             self.assertEqual(chant.segment_id, segment_1.id)
         for chant in source_2_chants:
             self.assertEqual(chant.segment_id, segment_2.id)
+        for sequence in sequence_source_sequences:
+            self.assertEqual(sequence.segment_id, segment_2.id)