3b1b-manim/manimlib/mobject/svg/mtex_mobject.py
2022-04-18 18:47:57 +08:00

360 lines
12 KiB
Python

from __future__ import annotations
import itertools as it
import re
from manimlib.mobject.svg.labelled_string import LabelledString
from manimlib.utils.tex_file_writing import display_during_execution
from manimlib.utils.tex_file_writing import get_tex_config
from manimlib.utils.tex_file_writing import tex_to_svg_file
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from colour import Color
from typing import Iterable, Union
from manimlib.mobject.types.vectorized_mobject import VGroup
ManimColor = Union[str, Color]
Span = tuple[int, int]
Selector = Union[
str,
re.Pattern,
tuple[Union[int, None], Union[int, None]],
Iterable[Union[
str,
re.Pattern,
tuple[Union[int, None], Union[int, None]]
]]
]
SCALE_FACTOR_PER_FONT_POINT = 0.001
TEX_COLOR_COMMANDS_DICT = {
"\\color": (1, False),
"\\textcolor": (1, False),
"\\pagecolor": (1, True),
"\\colorbox": (1, True),
"\\fcolorbox": (2, True),
}
class MTex(LabelledString):
CONFIG = {
"font_size": 48,
"alignment": "\\centering",
"tex_environment": "align*",
"tex_to_color_map": {},
}
def __init__(self, tex_string: str, **kwargs):
# Prevent from passing an empty string.
if not tex_string.strip():
tex_string = "\\\\"
self.tex_string = tex_string
super().__init__(tex_string, **kwargs)
self.set_color_by_tex_to_color_map(self.tex_to_color_map)
self.scale(SCALE_FACTOR_PER_FONT_POINT * self.font_size)
@property
def hash_seed(self) -> tuple:
return (
self.__class__.__name__,
self.svg_default,
self.path_string_config,
self.base_color,
self.isolate,
self.tex_string,
self.alignment,
self.tex_environment,
self.tex_to_color_map
)
def get_file_path_by_content(self, content: str) -> str:
tex_config = get_tex_config()
full_tex = tex_config["tex_body"].replace(
tex_config["text_to_replace"],
content
)
with display_during_execution(f"Writing \"{self.string}\""):
file_path = tex_to_svg_file(full_tex)
return file_path
def parse(self) -> None:
self.backslash_indices = self.get_backslash_indices()
self.command_spans = self.get_command_spans()
self.brace_spans = self.get_brace_spans()
self.script_char_indices = self.get_script_char_indices()
self.script_content_spans = self.get_script_content_spans()
self.script_spans = self.get_script_spans()
self.command_repl_items = self.get_command_repl_items()
super().parse()
# Toolkits
@staticmethod
def get_color_command_str(rgb_int: int) -> str:
rg, b = divmod(rgb_int, 256)
r, g = divmod(rg, 256)
return f"\\color[RGB]{{{r}, {g}, {b}}}"
# Parsing
def get_backslash_indices(self) -> list[int]:
# The latter of `\\` doesn't count.
return self.find_indices(r"\\.")
def get_command_spans(self) -> list[Span]:
return [
self.match(r"\\(?:[a-zA-Z]+|.)", pos=index).span()
for index in self.backslash_indices
]
def get_unescaped_char_indices(self, char: str) -> list[int]:
return list(filter(
lambda index: index - 1 not in self.backslash_indices,
self.find_indices(re.escape(char))
))
def get_brace_spans(self) -> list[Span]:
span_begins = []
span_ends = []
span_begins_stack = []
char_items = sorted([
(index, char)
for char in "{}"
for index in self.get_unescaped_char_indices(char)
])
for index, char in char_items:
if char == "{":
span_begins_stack.append(index)
else:
if not span_begins_stack:
raise ValueError("Missing '{' inserted")
span_begins.append(span_begins_stack.pop())
span_ends.append(index + 1)
if span_begins_stack:
raise ValueError("Missing '}' inserted")
return list(zip(span_begins, span_ends))
def get_script_char_indices(self) -> list[int]:
return list(it.chain(*[
self.get_unescaped_char_indices(char)
for char in "_^"
]))
def get_script_content_spans(self) -> list[Span]:
result = []
script_entity_dict = dict(it.chain(
self.brace_spans,
self.command_spans
))
for index in self.script_char_indices:
span_begin = self.match(r"\s*", pos=index + 1).end()
if span_begin in script_entity_dict.keys():
span_end = script_entity_dict[span_begin]
else:
match_obj = self.match(r".", pos=span_begin)
if match_obj is None:
continue
span_end = match_obj.end()
result.append((span_begin, span_end))
return result
def get_script_spans(self) -> list[Span]:
return [
(
self.match(r"[\s\S]*?(\s*)$", endpos=index).start(1),
script_content_span[1]
)
for index, script_content_span in zip(
self.script_char_indices, self.script_content_spans
)
]
def get_command_repl_items(self) -> list[tuple[Span, str]]:
result = []
brace_spans_dict = dict(self.brace_spans)
brace_begins = list(brace_spans_dict.keys())
for cmd_span in self.command_spans:
cmd_name = self.get_substr(cmd_span)
if cmd_name not in TEX_COLOR_COMMANDS_DICT.keys():
continue
n_braces, substitute_cmd = TEX_COLOR_COMMANDS_DICT[cmd_name]
span_begin, span_end = cmd_span
for _ in range(n_braces):
span_end = brace_spans_dict[min(filter(
lambda index: index >= span_end,
brace_begins
))]
if substitute_cmd:
repl_str = cmd_name + n_braces * "{black}"
else:
repl_str = ""
result.append(((span_begin, span_end), repl_str))
return result
def get_skippable_indices(self) -> list[int]:
return list(it.chain(
self.find_indices(r"\s"),
self.script_char_indices
))
def get_entity_spans(self) -> list[Span]:
return self.command_spans.copy()
def get_bracket_spans(self) -> list[Span]:
return self.brace_spans.copy()
def get_extra_isolated_items(self) -> list[tuple[Span, dict[str, str]]]:
result = []
# Match paired double braces (`{{...}}`).
reversed_brace_spans_dict = dict(sorted([
pair[::-1] for pair in self.brace_spans
]))
skip = False
for prev_brace_end, brace_end in self.get_neighbouring_pairs(
list(reversed_brace_spans_dict.keys())
):
if skip:
skip = False
continue
if brace_end != prev_brace_end + 1:
continue
brace_begin = reversed_brace_spans_dict[brace_end]
prev_brace_begin = reversed_brace_spans_dict[prev_brace_end]
if brace_begin != prev_brace_begin - 1:
continue
result.append((brace_begin, brace_end))
skip = True
result.extend(it.chain(*[
self.find_spans_by_selector(selector)
for selector in self.tex_to_color_map.keys()
]))
return [(span, {}) for span in result]
def get_label_span_list(self) -> list[Span]:
result = self.script_content_spans.copy()
reversed_script_spans_dict = dict([
script_span[::-1] for script_span in self.script_spans
])
for span_begin, span_end in self.specified_spans:
while span_end in reversed_script_spans_dict.keys():
span_end = reversed_script_spans_dict[span_end]
if span_begin >= span_end:
continue
shrinked_span = (span_begin, span_end)
if shrinked_span in result:
continue
result.append(shrinked_span)
return result
def get_content(self, is_labelled: bool) -> str:
if is_labelled:
extended_label_span_list = []
script_spans_dict = dict(self.script_spans)
for span in self.label_span_list:
if span not in self.script_content_spans:
span_begin, span_end = span
while span_end in script_spans_dict.keys():
span_end = script_spans_dict[span_end]
span = (span_begin, span_end)
extended_label_span_list.append(span)
inserted_string_pairs = [
(span, (
"{{" + self.get_color_command_str(label + 1),
"}}"
))
for label, span in enumerate(extended_label_span_list)
]
span_repl_dict = self.generate_span_repl_dict(
inserted_string_pairs, self.command_repl_items
)
else:
span_repl_dict = {}
result = self.get_replaced_substr(self.full_span, span_repl_dict)
if self.tex_environment:
if isinstance(self.tex_environment, str):
prefix = f"\\begin{{{self.tex_environment}}}"
suffix = f"\\end{{{self.tex_environment}}}"
else:
prefix, suffix = self.tex_environment
result = "\n".join([prefix, result, suffix])
if self.alignment:
result = "\n".join([self.alignment, result])
if not is_labelled:
result = "\n".join([
self.get_color_command_str(self.base_color_int),
result
])
return result
# Selector
def get_cleaned_substr(self, span: Span) -> str:
if not self.brace_spans:
brace_begins, brace_ends = [], []
else:
brace_begins, brace_ends = zip(*self.brace_spans)
left_brace_indices = list(brace_begins)
right_brace_indices = [index - 1 for index in brace_ends]
skippable_indices = list(it.chain(
self.skippable_indices,
left_brace_indices,
right_brace_indices
))
shrinked_span = self.shrink_span(span, skippable_indices)
if shrinked_span[0] >= shrinked_span[1]:
return ""
# Balance braces.
unclosed_left_braces = 0
unclosed_right_braces = 0
for index in range(*shrinked_span):
if index in left_brace_indices:
unclosed_left_braces += 1
elif index in right_brace_indices:
if unclosed_left_braces == 0:
unclosed_right_braces += 1
else:
unclosed_left_braces -= 1
return "".join([
unclosed_right_braces * "{",
self.get_substr(shrinked_span),
unclosed_left_braces * "}"
])
# Method alias
def get_parts_by_tex(self, selector: Selector, **kwargs) -> VGroup:
return self.select_parts(selector, **kwargs)
def get_part_by_tex(self, selector: Selector, **kwargs) -> VGroup:
return self.select_part(selector, **kwargs)
def set_color_by_tex(
self, selector: Selector, color: ManimColor, **kwargs
):
return self.set_parts_color(selector, color, **kwargs)
def set_color_by_tex_to_color_map(
self, color_map: dict[Selector, ManimColor], **kwargs
):
return self.set_parts_color_by_dict(color_map, **kwargs)
def get_tex(self) -> str:
return self.get_string()
class MTexText(MTex):
CONFIG = {
"tex_environment": None,
}