Enhance create_anatomy_questions management command with optional image generation and improved argument handling

.gitignore

@@ -1,4 +1,5 @@
 media/
+mediafiles/
 
 __pycache__
 .vscode/

anatomy/management/commands/create_anatomy_questions.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 import os
+import io
 from uuid import uuid4
-from typing import List
+from typing import List, Tuple
 
 from django.core.management.base import BaseCommand, CommandError
 from django.core.files import File
 from django.contrib.auth import get_user_model
 
-from anatomy.models import AnatomyQuestion, Exam, ExamQuestionDetail, QuestionType, Modality, Structure
+from PIL import Image
+
+from anatomy.models import AnatomyQuestion, Exam, ExamQuestionDetail, QuestionType, Structure
 from generic.models import Modality as GenericModality  # type: ignore
 
 
@@ -18,9 +21,13 @@ class Command(BaseCommand):
     )
 
     def add_arguments(self, parser):
-        parser.add_argument("--exam-id", type=int, required=True, help="ID of the anatomy.Exam to assign questions to")
+        parser.add_argument("--exam-id", type=int, required=False, help="ID of the anatomy.Exam to assign questions to (if omitted a new exam will be created)")
+        parser.add_argument("--exam-name", type=str, default=None, help="Name to use when creating a new exam if --exam-id omitted")
         parser.add_argument("--count", type=int, default=1, help="Number of questions to create")
-        parser.add_argument("--image-path", type=str, required=True, help="Path to image file to use for each created question")
+        parser.add_argument("--image-path", type=str, required=False, help="Path to image file to use for each created question")
+        parser.add_argument("--generate-image", action="store_true", help="Generate a placeholder image automatically if --image-path is not provided or not found")
+        parser.add_argument("--image-size", type=str, default="800x600", help="Generated image size WIDTHxHEIGHT, e.g. 800x600")
+        parser.add_argument("--image-color", type=str, default="#777777", help="Background color for generated image (hex) e.g. #ffffff")
         parser.add_argument("--description", type=str, default="Generated question {i}", help="Description template; use {i} for index")
         parser.add_argument("--question-type-id", type=int, default=None, help="Optional QuestionType id to assign")
         parser.add_argument("--modality-id", type=int, default=None, help="Optional Modality id to assign (generic.Modality id)")
@@ -29,9 +36,12 @@ class Command(BaseCommand):
         parser.add_argument("--start-sort", type=int, default=1, help="Starting sort_order for the first question")
 
     def handle(self, *args, **options):
-        exam_id = options["exam_id"]
+        exam_id = options.get("exam_id")
         count = options["count"]
-        image_path = options["image_path"]
+        image_path = options.get("image_path")
+        generate_image = options.get("generate_image")
+        image_size_str = options.get("image_size") or "800x600"
+        image_color = options.get("image_color") or "#777777"
         description_template = options["description"]
         qtype_id = options.get("question_type_id")
         modality_id = options.get("modality_id")
@@ -39,13 +49,25 @@ class Command(BaseCommand):
         author_ids: List[int] = options.get("author_id") or []
         sort = options.get("start_sort") or 1
 
+        exam = None
+        if exam_id:
             try:
                 exam = Exam.objects.get(pk=exam_id)
             except Exam.DoesNotExist:
                 raise CommandError(f"Exam with id {exam_id} does not exist")
+        else:
+            # create new exam
+            exam_name = options.get("exam_name") or f"Auto Anatomy Exam {uuid4().hex[:8]}"
+            exam = Exam(name=exam_name, active=True, open_access=True)
+            exam.save()
+            self.stdout.write(self.style.SUCCESS(f"Created exam id={exam.id} name='{exam.name}'"))
 
-        if not os.path.exists(image_path):
+        # validate image path if provided
+        if image_path and not os.path.exists(image_path):
+            if not generate_image:
                 raise CommandError(f"Image file not found at path: {image_path}")
+            else:
+                image_path = None
 
         # Resolve optional foreign keys
         qtype = None
@@ -82,17 +104,50 @@ class Command(BaseCommand):
                 raise CommandError(f"Author user id {uid} not found")
             authors.append(u)
 
+        # parse image size
+        try:
+            w, h = image_size_str.lower().split("x")
+            image_size: Tuple[int, int] = (int(w), int(h))
+        except Exception:
+            raise CommandError("Invalid --image-size; expected WIDTHxHEIGHT, e.g. 800x600")
+
         created = []
 
         for i in range(1, count + 1):
             desc = description_template.format(i=i)
 
-            # Prepare a unique filename to avoid collisions
+            # determine image source: either provided path or generated image
-            basename = os.path.basename(image_path)
+            unique_name = f"{uuid4().hex}_image.jpg"
-            unique_name = f"{uuid4().hex}_{basename}"
 
+            if image_path:
+                # use provided file
                 with open(image_path, "rb") as f:
                     django_file = File(f, name=unique_name)
+                    img_file_for_save = django_file
+                    # create and save question
+                    q = AnatomyQuestion(
+                        question_type=qtype,
+                        description=desc,
+                        answer_help="",
+                    )
+
+                    if modality:
+                        q.modality = modality
+
+                    if structure:
+                        q.structure = structure
+
+                    q.image.save(unique_name, img_file_for_save, save=False)
+                    q.save()
+
+            else:
+                # generate an image in memory
+                buf = io.BytesIO()
+                img = Image.new("RGB", image_size, image_color)
+                img.save(buf, format="JPEG")
+                buf.seek(0)
+                django_file = File(buf, name=unique_name)
+
                 q = AnatomyQuestion(
                     question_type=qtype,
                     description=desc,
@@ -105,7 +160,6 @@ class Command(BaseCommand):
                 if structure:
                     q.structure = structure
 
-                # Assign image then save
                 q.image.save(unique_name, django_file, save=False)
                 q.save()
 
@@ -114,11 +168,9 @@ class Command(BaseCommand):
                 q.author.set(authors)
 
             # add to exam via the many-to-many through relationship
-                # use through_defaults for sort_order when possible
             try:
                 exam.exam_questions.add(q, through_defaults={"sort_order": sort})
             except TypeError:
-                    # older Django or in case add doesn't support through_defaults, create through model
                 ExamQuestionDetail.objects.create(exam=exam, anatomyquestion=q, sort_order=sort)
 
             self.stdout.write(self.style.SUCCESS(f"Created question id={q.id} and added to exam id={exam.id} with sort_order={sort}"))