manim/manimlib/mobject/svg/mtex_mobject.py

import itertools as it
import re

from manimlib.mobject.svg.svg_mobject import SVGMobject
from manimlib.mobject.types.vectorized_mobject import VMobject
from manimlib.mobject.types.vectorized_mobject import VGroup
from manimlib.utils.iterables import adjacent_pairs
from manimlib.utils.iterables import remove_list_redundancies
from manimlib.utils.tex_file_writing import tex_to_svg_file
from manimlib.utils.tex_file_writing import get_tex_config
from manimlib.utils.tex_file_writing import display_during_execution


SCALE_FACTOR_PER_FONT_POINT = 0.001


tex_hash_to_mob_map = {}


class _LabelledTex(SVGMobject):
    CONFIG = {
        "height": None,
        "path_string_config": {
            "should_subdivide_sharp_curves": True,
            "should_remove_null_curves": True,
        },
    }

    def get_mobjects_from(self, element):
        result = super().get_mobjects_from(element)
        for mob in result:
            if not hasattr(mob, "label_str"):
                mob.label_str = ""
        try:
            label_str = element.getAttribute("fill")
            if label_str:
                if len(label_str) == 4:
                    # "#RGB" => "#RRGGBB"
                    label_str = "#" + "".join([c * 2 for c in label_str[1:]])
                for mob in result:
                    mob.label_str = label_str
        except:
            pass
        return result


class _TexSpan(object):
    def __init__(self, script_type, label, referring_labels):
        # 0 for normal, 1 for subscript, 2 for superscript.
        self.script_type = script_type
        self.label = label
        self.referring_labels = referring_labels


class MTex(VMobject):
    CONFIG = {
        "fill_opacity": 1.0,
        "stroke_width": 0,
        "should_center": True,
        "font_size": 48,
        "height": None,
        "organize_left_to_right": False,
        "alignment": "\\centering",
        "math_mode": True,
        "isolate": [],
        "tex_to_color_map": {},
    }

    def __init__(self, tex_string, **kwargs):
        super().__init__(**kwargs)
        self.tex_string = tex_string
        self.parse_tex()

        full_tex = self.get_tex_file_body()
        hash_val = hash(full_tex)
        if hash_val not in tex_hash_to_mob_map:
            with display_during_execution(f"Writing \"{tex_string}\""):
                filename = tex_to_svg_file(full_tex)
                svg_mob = _LabelledTex(filename)
                tex_hash_to_mob_map[hash_val] = svg_mob
        self.add(*[
            submob.copy()
            for submob in tex_hash_to_mob_map[hash_val]
        ])
        self.build_structure()

        self.init_colors()
        self.set_color_by_tex_to_color_map(self.tex_to_color_map)

        if self.height is None:
            self.scale(SCALE_FACTOR_PER_FONT_POINT * self.font_size)
        if self.organize_left_to_right:
            self.organize_submobjects_left_to_right()

    def add_tex_span(self, span_tuple, script_type=0, referring_labels=None):
        if referring_labels is None:
            # Should be additionally labelled.
            label = self.current_label
            self.current_label += 1
            referring_labels = [label]
        else:
            label = -1

        # 0 for normal, 1 for subscript, 2 for superscript.
        # Only those spans with `label != -1` will be colored.
        tex_span = _TexSpan(script_type, label, referring_labels)
        self.tex_spans_dict[span_tuple] = tex_span

    def parse_tex(self):
        self.tex_spans_dict = {}
        self.current_label = 0
        self.break_up_by_braces()
        self.break_up_by_scripts()
        self.break_up_by_additional_strings()
        self.analyse_referring_colors()

    def break_up_by_braces(self):
        span_tuples = []
        left_brace_indices = []
        for match_obj in re.finditer(r"(?<!\\)(\{|\})", self.tex_string):
            if match_obj.group() == "{":
                left_brace_index = match_obj.span()[0]
                left_brace_indices.append(left_brace_index)
            else:
                left_brace_index = left_brace_indices.pop()
                right_brace_index = match_obj.span()[1]
                span_tuples.append((left_brace_index, right_brace_index))
        if left_brace_indices:
            self.raise_tex_parsing_error()

        for span_tuple in span_tuples:
            self.add_tex_span(span_tuple)

    def break_up_by_scripts(self):
        tex_string = self.tex_string
        brace_indices_dict = dict([
            (span_tuple[0], span_tuple)
            for span_tuple in self.tex_spans_dict.keys()
        ])
        for match_obj in re.finditer(r"((?<!\\)(_|\^)\s*)|(\s+(_|\^)\s*)", tex_string):
            script_type = 1 if "_" in match_obj.group() else 2
            token_begin, token_end = match_obj.span()
            if token_end in brace_indices_dict:
                content_span = brace_indices_dict[token_end]
            else:
                content_match_obj = re.match(r"\w|\\[a-zA-Z]+", tex_string[token_end:])
                if not content_match_obj:
                    self.raise_tex_parsing_error()
                content_span = tuple([
                    index + token_end for index in content_match_obj.span()
                ])
                self.add_tex_span(content_span)
            label = self.tex_spans_dict[content_span].label
            self.add_tex_span(
                (token_begin, content_span[1]),
                script_type=script_type,
                referring_labels=[label]
            )

    def break_up_by_additional_strings(self):
        additional_strings_to_break_up = remove_list_redundancies([
            *self.isolate, *self.tex_to_color_map.keys()
        ])
        if not additional_strings_to_break_up:
            return

        tex_string = self.tex_string
        all_span_tuples = list(self.tex_spans_dict.keys())
        for string in additional_strings_to_break_up:
            # Only matches non-overlapping strings.
            for match_obj in re.finditer(re.escape(string), tex_string):
                all_span_tuples.append(match_obj.span())

        # Deconstruct spans with subscripts & superscripts.
        script_spans_dict = dict([
            span_tuple[::-1]
            for span_tuple, tex_span in self.tex_spans_dict.items()
            if tex_span.script_type != 0
        ])
        for span_begin, span_end in all_span_tuples:
            while span_end in script_spans_dict:
                span_end = script_spans_dict[span_end]
            if span_begin >= span_end:
                continue
            span_tuple = (span_begin, span_end)
            if span_tuple not in self.tex_spans_dict:
                self.add_tex_span(span_tuple)

    def analyse_referring_colors(self):
        all_span_tuples = list(self.tex_spans_dict.keys())
        if not all_span_tuples:
            return

        for i, span_0 in enumerate(all_span_tuples):
            for j, span_1 in enumerate(all_span_tuples):
                if i == j:
                    continue
                tex_span_0 = self.tex_spans_dict[span_0]
                tex_span_1 = self.tex_spans_dict[span_1]
                if tex_span_0.label == -1:
                    continue
                if span_0[0] <= span_1[0] and span_0[1] >= span_1[1]:
                    tex_span_0.referring_labels.append(tex_span_1.label)

    def raise_tex_parsing_error(self):
        raise ValueError(f"Failed to parse tex: \"{self.tex_string}\"")

    def get_tex_file_body(self):
        new_tex = self.get_modified_expression()
        if self.math_mode:
            new_tex = "\\begin{align*}\n" + new_tex + "\n\\end{align*}"

        new_tex = self.alignment + "\n" + new_tex

        tex_config = get_tex_config()
        return tex_config["tex_body"].replace(
            tex_config["text_to_replace"],
            new_tex
        )

    def get_modified_expression(self):
        tex_string = self.tex_string
        if not self.tex_spans_dict:
            return tex_string

        indices_with_labels = sorted([
            (index, i, tex_span.label)
            for span_tuple, tex_span in self.tex_spans_dict.items()
            for i, index in enumerate(span_tuple)
            if tex_span.label != -1
        ], key=lambda t: (t[0], 1 - t[1]))
        # Add one more item to ensure all the substrings are joined.
        indices_with_labels.append((
            len(tex_string), 0, -1
        ))

        result = tex_string[: indices_with_labels[0][0]]
        for index_0_with_label, index_1_with_label in list(adjacent_pairs(indices_with_labels))[:-1]:
            index, flag, label = index_0_with_label
            if flag == 0:
                color_tuple = MTex.label_to_color_tuple(label)
                result += "".join([
                    "{{",
                    "\\color[RGB]",
                    "{" + ",".join(map(str, color_tuple)) + "}"
                ])
            else:
                result += "}}"
            result += tex_string[index : index_1_with_label[0]]
        return result

    @staticmethod
    def label_to_color_tuple(n):
        # Get a unique color different from black,
        # or the svg file will not include the color information.
        return (
            (n + 1) // 256 // 256,
            (n + 1) // 256 % 256,
            (n + 1) % 256
        )

    @staticmethod
    def color_str_to_label(color):
        return int(color[1:], 16) - 1

    def build_structure(self):
        # Simply pack together adjacent mobjects with the same label.
        new_submobjects = []
        new_submobject_components = []
        current_label_str = ""
        for submob in self.submobjects:
            if submob.label_str == current_label_str:
                new_submobject_components.append(submob)
            else:
                if new_submobject_components:
                    new_submobjects.append(VGroup(*new_submobject_components))
                new_submobject_components = [submob]
                current_label_str = submob.label_str
        if new_submobject_components:
            new_submobjects.append(VGroup(*new_submobject_components))

        self.set_submobjects(new_submobjects)
        return self

    def get_parts_by_tex(self, tex):
        result = VGroup()
        d = dict(self.tex_spans_dict.keys())
        for match_obj in re.finditer(re.escape(tex), self.tex_string):
            labels = []
            span_begin, span_end = match_obj.span()
            while span_begin < span_end and span_begin in d:
                next_span_begin = d[span_begin]
                referring_labels = self.tex_spans_dict[(span_begin, next_span_begin)].referring_labels
                labels.extend(referring_labels)
                span_begin = next_span_begin
            if span_begin != span_end:
                raise ValueError(f"Failed to get span of tex: \"{tex}\"")

            mob = VGroup(*filter(
                lambda submob: submob.label_str and MTex.color_str_to_label(submob.label_str) in labels,
                it.chain(*self.submobjects)
            ))
            result.add(mob)
        return result

    def get_part_by_tex(self, tex):
        all_parts = self.get_parts_by_tex(tex)
        return all_parts[0] if all_parts else None

    def set_color_by_tex(self, tex, color):
        self.get_parts_by_tex(tex).set_color(color)
        return self

    def set_color_by_tex_to_color_map(self, tex_to_color_map):
        for tex, color in list(tex_to_color_map.items()):
            self.set_color_by_tex(tex, color)
        return self


class MTexText(MTex):
    CONFIG = {
        "math_mode": False,
    }