prompt_user_for_choice now prompts for Scene Name or Number

manimlib/extract_scene.py

@@ -2,6 +2,8 @@ import inspect
 import sys
 import logging
 
+from numpy.lib.shape_base import split
+
 from manimlib.scene.scene import Scene
 from manimlib.config import get_custom_defaults
 
@@ -26,18 +28,20 @@ def is_child_scene(obj, module):
 
 def prompt_user_for_choice(scene_classes):
     name_to_class = {}
-    for scene_class in scene_classes:
+    max_digits = len(str(len(scene_classes)))
+    for idx, scene_class in enumerate(scene_classes, start=1):
         name = scene_class.__name__
-        print(name)
+        print(f"{str(idx).zfill(max_digits)}: {name}")
         name_to_class[name] = scene_class
     try:
         user_input = input(
-            "\nThat module has multiple scenes, which "
+            "\nThat module has multiple scenes, "
-            "ones would you like to render?\n Scene Name: "
+            "which ones would you like to render?"
+            "\nScene Name or Number: "
         )
         return [
-            name_to_class[user_input]
+            name_to_class[split_str] if not split_str.isnumeric() else scene_classes[int(split_str)-1]
-            for num_str in user_input.replace(" ", "").split(",")
+            for split_str in user_input.replace(" ", "").split(",")
         ]
     except KeyError:
         logging.log(logging.ERROR, "Invalid scene")