Fix bug when handling multi-line tex

manimlib/mobject/svg/labelled_string.py

@@ -57,6 +57,7 @@ class LabelledString(_StringSVG):
         self.pre_parse()
         self.parse()
         super().__init__(**kwargs)
+        self.post_parse()
 
     def get_file_path(self) -> str:
         return self.get_file_path_(use_plain_file=False)
@@ -108,6 +109,20 @@ class LabelledString(_StringSVG):
         self.label_span_list = self.get_label_span_list()
         self.check_overlapping()
 
+    def post_parse(self) -> None:
+        self.labelled_submobject_items = [
+            (submob.label, submob)
+            for submob in self.submobjects
+        ]
+        self.labelled_submobjects = self.get_labelled_submobjects()
+        self.specified_substrs = self.get_specified_substrs()
+        self.group_items = self.get_group_items()
+        self.group_substrs = self.get_group_substrs()
+        self.submob_groups = self.get_submob_groups()
+
+    def copy(self):
+        return self.deepcopy()
+
     # Toolkits
 
     def get_substr(self, span: Span) -> str:
@@ -118,10 +133,14 @@ class LabelledString(_StringSVG):
     ) -> Iterable[re.Match]:
         return re.compile(pattern, flags).finditer(self.string, **kwargs)
 
-    def search(self, pattern: str, flags: int = 0, **kwargs) -> re.Match:
+    def search(
+        self, pattern: str, flags: int = 0, **kwargs
+    ) -> re.Match | None:
         return re.compile(pattern, flags).search(self.string, **kwargs)
 
-    def match(self, pattern: str, flags: int = 0, **kwargs) -> re.Match:
+    def match(
+        self, pattern: str, flags: int = 0, **kwargs
+    ) -> re.Match | None:
         return re.compile(pattern, flags).match(self.string, **kwargs)
 
     def find_spans(self, pattern: str, **kwargs) -> list[Span]:
@@ -275,9 +294,7 @@ class LabelledString(_StringSVG):
     def color_to_label(color: ManimColor) -> int:
         rgb_tuple = color_to_int_rgb(color)
         rgb = LabelledString.rgb_to_int(rgb_tuple)
-        if rgb == 16777215:  # white
-            return -1
-        return rgb
+        return rgb - 1
 
     @abstractmethod
     def get_begin_color_command_str(int_rgb: int) -> str:
@@ -321,12 +338,11 @@ class LabelledString(_StringSVG):
         return []
 
     def get_specified_spans(self) -> list[Span]:
-        spans = [
-            self.full_span,
-            *self.internal_specified_spans,
-            *self.external_specified_spans,
-            *self.find_substrs(self.isolate)
-        ]
+        spans = list(it.chain(
+            self.internal_specified_spans,
+            self.external_specified_spans,
+            self.find_substrs(self.isolate)
+        ))
         shrinked_spans = list(filter(
             lambda span: span[0] < span[1],
             [self.shrink_span(span) for span in spans]
@@ -381,6 +397,9 @@ class LabelledString(_StringSVG):
 
     # Post-parsing
 
+    def get_labelled_submobjects(self) -> list[VMobject]:
+        return [submob for _, submob in self.labelled_submobject_items]
+
     def get_cleaned_substr(self, span: Span) -> str:
         span_repl_dict = dict.fromkeys(self.command_spans, "")
         return self.get_replaced_substr(span, span_repl_dict)
@@ -391,17 +410,14 @@ class LabelledString(_StringSVG):
             for span in self.specified_spans
         ])
 
-    def get_group_span_items(self) -> tuple[list[int], list[Span]]:
-        submob_labels = [submob.label for submob in self.submobjects]
-        if not submob_labels:
-            return [], []
-        return tuple(zip(*self.compress_neighbours(submob_labels)))
-
-    def get_group_substrs(self) -> list[str]:
-        group_labels, _ = self.get_group_span_items()
-        if not group_labels:
+    def get_group_items(self) -> list[tuple[str, VGroup]]:
+        if not self.labelled_submobject_items:
             return []
 
+        labels, labelled_submobjects = zip(*self.labelled_submobject_items)
+        group_labels, labelled_submob_spans = zip(
+            *self.compress_neighbours(labels)
+        )
         ordered_spans = [
             self.label_span_list[label] if label != -1 else self.full_span
             for label in group_labels
@@ -425,16 +441,27 @@ class LabelledString(_StringSVG):
                 interval_spans, (ordered_spans[0][0], ordered_spans[-1][1])
             )
         ]
-        return [
+        group_substrs = [
             self.get_cleaned_substr(span) if span[0] < span[1] else ""
             for span in shrinked_spans
         ]
+        submob_groups = VGroup(*[
+            VGroup(*labelled_submobjects[slice(*submob_span)])
+            for submob_span in labelled_submob_spans
+        ])
+        return list(zip(group_substrs, submob_groups))
 
-    def get_submob_groups(self) -> VGroup:
-        _, submob_spans = self.get_group_span_items()
+    def get_group_substrs(self) -> list[str]:
+        return [group_substr for group_substr, _ in self.group_items]
+
+    def get_submob_groups(self) -> list[VGroup]:
+        return [submob_group for _, submob_group in self.group_items]
+
+    def get_parts_by_group_substr(self, substr: str) -> VGroup:
         return VGroup(*[
-            VGroup(*self.submobjects[slice(*submob_span)])
-            for submob_span in submob_spans
+            group
+            for group_substr, group in self.group_items
+            if group_substr == substr
         ])
 
     # Selector
@@ -477,7 +504,7 @@ class LabelledString(_StringSVG):
             span_begin = next_begin
         return result
 
-    def get_parts_by_custom_span(self, custom_span: Span, **kwargs) -> VGroup:
+    def get_part_by_custom_span(self, custom_span: Span, **kwargs) -> VGroup:
         labels = [
             label for label, span in enumerate(self.label_span_list)
             if any([
@@ -487,10 +514,10 @@ class LabelledString(_StringSVG):
                 )
             ])
         ]
-        return VGroup(*filter(
-            lambda submob: submob.label in labels,
-            self.submobjects
-        ))
+        return VGroup(*[
+            submob for label, submob in self.labelled_submobject_items
+            if label in labels
+        ])
 
     def get_parts_by_string(
         self, substr: str, case_sensitive: bool = True, **kwargs
@@ -499,19 +526,10 @@ class LabelledString(_StringSVG):
         if not case_sensitive:
             flags |= re.I
         return VGroup(*[
-            self.get_parts_by_custom_span(span, **kwargs)
+            self.get_part_by_custom_span(span, **kwargs)
             for span in self.find_substr(substr, flags=flags)
         ])
 
-    def get_parts_by_group_substr(self, substr: str) -> VGroup:
-        return VGroup(*[
-            group
-            for group, group_substr in zip(
-                self.get_submob_groups(), self.get_group_substrs()
-            )
-            if group_substr == substr
-        ])
-
     def get_part_by_string(
         self, substr: str, index: int = 0, **kwargs
     ) -> VMobject:
@@ -528,13 +546,5 @@ class LabelledString(_StringSVG):
             self.set_color_by_string(substr, color, **kwargs)
         return self
 
-    def indices_of_part(self, part: Iterable[VMobject]) -> list[int]:
-        return [self.submobjects.index(submob) for submob in part]
-
-    def indices_lists_of_parts(
-        self, parts: Iterable[Iterable[VMobject]]
-    ) -> list[list[int]]:
-        return [self.indices_of_part(part) for part in parts]
-
     def get_string(self) -> str:
         return self.string