From 2536b7fafbee98cc01f6388d95221f8375f3112f Mon Sep 17 00:00:00 2001
From: Christopher Hampson <christopher.hampson.1989@gmail.com>
Date: Sun, 1 Sep 2024 16:45:30 +0100
Subject: [PATCH] feat: preserve original TeX colors (#3903)

* feat: preserve original Tex colors

* added test for preserving tex color

* fix: fixed breaking change to tex color inheritance

* updated docstring and comment

* Update manim/mobject/text/tex_mobject.py

Co-authored-by: Aarush Deshpande <110117391+JasonGrace2282@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Aarush Deshpande <110117391+JasonGrace2282@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 manim/mobject/text/tex_mobject.py             | 28 +++++++++++------
 .../test_graphical_units/test_tex_mobject.py  | 30 +++++++++++++++++++
 2 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/manim/mobject/text/tex_mobject.py b/manim/mobject/text/tex_mobject.py
index e8c09bbac1..26334a60d9 100644
--- a/manim/mobject/text/tex_mobject.py
+++ b/manim/mobject/text/tex_mobject.py
@@ -12,7 +12,7 @@
 
 from __future__ import annotations
 
-from manim.utils.color import ManimColor
+from manim.utils.color import BLACK, ManimColor, ParsableManimColor
 
 __all__ = [
     "SingleStringMathTex",
@@ -62,12 +62,11 @@ def __init__(
         tex_environment: str = "align*",
         tex_template: TexTemplate | None = None,
         font_size: float = DEFAULT_FONT_SIZE,
+        color: ParsableManimColor | None = None,
         **kwargs,
     ):
-        if kwargs.get("color") is None:
-            # makes it so that color isn't explicitly passed for these mobs,
-            # and can instead inherit from the parent
-            kwargs["color"] = VMobject().color
+        if color is None:
+            color = VMobject().color
 
         self._font_size = font_size
         self.organize_left_to_right = organize_left_to_right
@@ -88,6 +87,7 @@ def __init__(
             should_center=should_center,
             stroke_width=stroke_width,
             height=height,
+            color=color,
             path_string_config={
                 "should_subdivide_sharp_curves": True,
                 "should_remove_null_curves": True,
@@ -210,10 +210,16 @@ def get_tex_string(self):
         return self.tex_string
 
     def init_colors(self, propagate_colors=True):
-        if config.renderer == RendererType.OPENGL:
-            super().init_colors()
-        elif config.renderer == RendererType.CAIRO:
-            super().init_colors(propagate_colors=propagate_colors)
+        for submobject in self.submobjects:
+            # needed to preserve original (non-black)
+            # TeX colors of individual submobjects
+            if submobject.color != BLACK:
+                continue
+            submobject.color = self.color
+            if config.renderer == RendererType.OPENGL:
+                submobject.init_colors()
+            elif config.renderer == RendererType.CAIRO:
+                submobject.init_colors(propagate_colors=propagate_colors)
 
 
 class MathTex(SingleStringMathTex):
@@ -426,6 +432,10 @@ def sort_alphabetically(self):
 class Tex(MathTex):
     r"""A string compiled with LaTeX in normal mode.
 
+    The color can be set using
+    the ``color`` argument. Any parts of the ``tex_string`` that are colored by the
+    TeX commands ``\color`` or ``\textcolor`` will retain their original color.
+
     Tests
     -----
 
diff --git a/tests/test_graphical_units/test_tex_mobject.py b/tests/test_graphical_units/test_tex_mobject.py
index 3286352a8e..aabc8e01f9 100644
--- a/tests/test_graphical_units/test_tex_mobject.py
+++ b/tests/test_graphical_units/test_tex_mobject.py
@@ -24,3 +24,33 @@ def test_set_opacity_by_tex(scene):
     tex = MathTex("f(x) = y", substrings_to_isolate=["f(x)"])
     tex.set_opacity_by_tex("f(x)", 0.2, 0.5)
     scene.add(tex)
+
+
+def test_preserve_tex_color():
+    """Test that Tex preserves original tex colors."""
+    template = TexTemplate(preamble=r"\usepackage{xcolor}")
+    Tex.set_default(tex_template=template)
+
+    txt = Tex(r"\textcolor{red}{Hello} World")
+    assert len(txt[0].submobjects) == 10
+    assert all(char.fill_color.to_hex() == "#FF0000" for char in txt[0][:5])  # "Hello"
+    assert all(
+        char.fill_color.to_hex() == WHITE.to_hex() for char in txt[0][-5:]
+    )  # "World"
+
+    txt = Tex(r"\textcolor{red}{Hello} World", color=BLUE)
+    assert len(txt[0].submobjects) == 10
+    assert all(char.fill_color.to_hex() == "#FF0000" for char in txt[0][:5])  # "Hello"
+    assert all(
+        char.fill_color.to_hex() == BLUE.to_hex() for char in txt[0][-5:]
+    )  # "World"
+
+    Tex.set_default(color=GREEN)
+    txt = Tex(r"\textcolor{red}{Hello} World")
+    assert len(txt[0].submobjects) == 10
+    assert all(char.fill_color.to_hex() == "#FF0000" for char in txt[0][:5])  # "Hello"
+    assert all(
+        char.fill_color.to_hex() == GREEN.to_hex() for char in txt[0][-5:]
+    )  # "World"
+
+    Tex.set_default()