Cache string coloring functions (#2977)

* Cache color functions until the color parameter is set

* Update pwndbg/color/theme.py

Co-authored-by: Disconnect3d <dominik.b.czarnota@gmail.com>

---------

Co-authored-by: Disconnect3d <dominik.b.czarnota@gmail.com>

pwndbg/color/__init__.py

@@ -164,7 +164,7 @@ class ColorParamSpec(NamedTuple):
 class ColorConfig:
     def __init__(self, namespace: str, params: List[ColorParamSpec]) -> None:
         self._namespace = namespace
-        self._params: Dict[str, Parameter] = {}
+        self._params: Dict[str, theme.ColorParameter] = {}
         for param in params:
             self._params[param.name] = theme.add_color_param(
                 f"{self._namespace}-{param.name}-color", param.default, param.doc
@@ -173,7 +173,7 @@ class ColorConfig:
     def __getattr__(self, attr: str) -> Callable[[str], str]:
         param_name = attr.replace("_", "-")
         if param_name in self._params:
-            return generateColorFunction(self._params[param_name])
+            return self._params[param_name].color_function
 
         raise AttributeError(f"ColorConfig object for {self._namespace} has no attribute '{attr}'")
 

pwndbg/color/context.py

@@ -2,8 +2,6 @@ from __future__ import annotations
 
 from typing import List
 
-from pwndbg import config
-from pwndbg.color import generateColorFunction
 from pwndbg.color import theme
 from pwndbg.lib.regs import BitFlags
 
@@ -43,51 +41,51 @@ config_comment = theme.add_color_param("comment-color", "gray", "color for comme
 
 
 def prefix(x: object) -> str:
-    return generateColorFunction(config.code_prefix_color)(x)
+    return config_prefix_color.color_function(x)
 
 
 def highlight(x: object) -> str:
-    return generateColorFunction(config.highlight_color)(x)
+    return config_highlight_color.color_function(x)
 
 
 def register(x: object) -> str:
-    return generateColorFunction(config.context_register_color)(x)
+    return config_register_color.color_function(x)
 
 
 def register_changed(x: object) -> str:
-    return generateColorFunction(config.context_register_changed_color)(x)
+    return config_register_changed_color.color_function(x)
 
 
 def flag_bracket(x: object) -> str:
-    return generateColorFunction(config.context_flag_bracket_color)(x)
+    return config_flag_bracket_color.color_function(x)
 
 
 def flag_value(x: object) -> str:
-    return generateColorFunction(config.context_flag_value_color)(x)
+    return config_flag_value_color.color_function(x)
 
 
 def flag_set(x: object) -> str:
-    return generateColorFunction(config.context_flag_set_color)(x)
+    return config_flag_set_color.color_function(x)
 
 
 def flag_unset(x: object) -> str:
-    return generateColorFunction(config.context_flag_unset_color)(x)
+    return config_flag_unset_color.color_function(x)
 
 
 def flag_changed(x: object) -> str:
-    return generateColorFunction(config.context_flag_changed_color)(x)
+    return config_flag_changed_color.color_function(x)
 
 
 def banner(x: object) -> str:
-    return generateColorFunction(config.banner_color)(x)
+    return config_banner_color.color_function(x)
 
 
 def banner_title(x: object) -> str:
-    return generateColorFunction(config.banner_title_color)(x)
+    return config_banner_title.color_function(x)
 
 
 def comment(x: object) -> str:
-    return generateColorFunction(config.comment_color)(x)
+    return config_comment.color_function(x)
 
 
 def format_flags(value: int | None, flags: BitFlags, last: int | None = None):

pwndbg/color/enhance.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from pwndbg import config
-from pwndbg.color import generateColorFunction
 from pwndbg.color import theme
 
 config_integer_color = theme.add_color_param(
@@ -19,16 +17,16 @@ config_unknown_color = theme.add_color_param(
 
 
 def integer(x):
-    return generateColorFunction(config.enhance_integer_value_color)(x)
+    return config_integer_color.color_function(x)
 
 
 def string(x):
-    return generateColorFunction(config.enhance_string_value_color)(x)
+    return config_string_color.color_function(x)
 
 
 def comment(x):
-    return generateColorFunction(config.enhance_comment_color)(x)
+    return config_comment_color.color_function(x)
 
 
 def unknown(x):
-    return generateColorFunction(config.enhance_unknown_color)(x)
+    return config_unknown_color.color_function(x)

pwndbg/color/hexdump.py

@@ -34,31 +34,31 @@ config_highlight_group_lsb = theme.add_param(
 
 
 def normal(x: str) -> str:
-    return generateColorFunction(config.hexdump_normal_color)(x)
+    return config_normal.color_function(x)
 
 
 def printable(x: str) -> str:
-    return generateColorFunction(config.hexdump_printable_color)(x)
+    return config_printable.color_function(x)
 
 
 def zero(x: str) -> str:
-    return generateColorFunction(config.hexdump_zero_color)(x)
+    return config_zero.color_function(x)
 
 
 def special(x: str) -> str:
-    return generateColorFunction(config.hexdump_special_color)(x)
+    return config_special.color_function(x)
 
 
 def offset(x: str) -> str:
-    return generateColorFunction(config.hexdump_offset_color)(x)
+    return config_offset.color_function(x)
 
 
 def address(x: str) -> str:
-    return generateColorFunction(config.hexdump_address_color)(x)
+    return config_address.color_function(x)
 
 
 def separator(x: str) -> str:
-    return generateColorFunction(config.hexdump_separator_color)(x)
+    return config_separator.color_function(x)
 
 
 def highlight_group_lsb(x: str) -> str:

pwndbg/color/message.py

@@ -2,9 +2,6 @@ from __future__ import annotations
 
 from typing import Callable
 
-import pwndbg.lib.config
-from pwndbg import config
-from pwndbg.color import generateColorFunction
 from pwndbg.color import theme
 
 config_status_on_color = theme.add_color_param(
@@ -41,72 +38,70 @@ config_signal_color = theme.add_color_param(
     "message-signal-color", "bold,red", "color of signal messages"
 )
 
-config_prompt_color: pwndbg.lib.config.Parameter = theme.add_color_param(
-    "prompt-color", "bold,red", "prompt color"
-)
-config_prompt_alive_color: pwndbg.lib.config.Parameter = theme.add_color_param(
+config_prompt_color = theme.add_color_param("prompt-color", "bold,red", "prompt color")
+config_prompt_alive_color = theme.add_color_param(
     "prompt-alive-color", "bold,green", "prompt alive color"
 )
 
 
 def on(msg: object) -> str:
-    return generateColorFunction(config.message_status_on_color)(msg)
+    return config_status_on_color.color_function(msg)
 
 
 def off(msg: object) -> str:
-    return generateColorFunction(config.message_status_off_color)(msg)
+    return config_status_off_color.color_function(msg)
 
 
 def notice(msg: object) -> str:
-    return generateColorFunction(config.message_notice_color)(msg)
+    return config_notice_color.color_function(msg)
 
 
 def hint(msg: object) -> str:
-    return generateColorFunction(config.message_hint_color)(msg)
+    return config_hint_color.color_function(msg)
 
 
 def success(msg: object) -> str:
-    return generateColorFunction(config.message_success_color)(msg)
+    return config_success_color.color_function(msg)
 
 
 def debug(msg: object) -> str:
-    return generateColorFunction(config.message_warning_color)(msg)
+    return config_debug_color.color_function(msg)
 
 
 def info(msg: object) -> str:
-    return generateColorFunction(config.message_warning_color)(msg)
+    return config_info_color.color_function(msg)
 
 
 def warn(msg: object) -> str:
-    return generateColorFunction(config.message_warning_color)(msg)
+    return config_warning_color.color_function(msg)
 
 
 def error(msg: object) -> str:
-    return generateColorFunction(config.message_error_color)(msg)
+    return config_error_color.color_function(msg)
 
 
 def system(msg: object) -> str:
-    return generateColorFunction(config.message_system_color)(msg)
+    return config_system_color.color_function(msg)
 
 
 def exit(msg: object) -> str:
-    return generateColorFunction(config.message_exit_color)(msg)
+    return config_exit_color.color_function(msg)
 
 
 def breakpoint(msg: object) -> str:
-    return generateColorFunction(config.message_breakpoint_color)(msg)
+    return config_breakpoint_color.color_function(msg)
 
 
 def signal(msg: object) -> str:
-    return generateColorFunction(config.message_signal_color)(msg)
+    return config_signal_color.color_function(msg)
 
 
 def prompt(msg: object) -> str:
-    return generateColorFunction(config.prompt_color)(msg)
+    return config_prompt_color.color_function(msg)
 
 
 def alive_prompt(msg: object) -> str:
-    return generateColorFunction(config.prompt_alive_color)(msg)
+    return config_prompt_alive_color.color_function(msg)
 
 
 def readline_escape(func_message: Callable[[str], str], text: str) -> str:

pwndbg/color/telescope.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from pwndbg import config
-from pwndbg.color import generateColorFunction
 from pwndbg.color import theme
 
 offset_color = theme.add_color_param(
@@ -28,20 +26,20 @@ repeating_marker_color = theme.add_color_param(
 
 
 def offset(x: object) -> str:
-    return generateColorFunction(config.telescope_offset_color)(x)
+    return offset_color.color_function(x)
 
 
 def register(x: object) -> str:
-    return generateColorFunction(config.telescope_register_color)(x)
+    return register_color.color_function(x)
 
 
 def separator(x: object) -> str:
-    return generateColorFunction(config.telescope_offset_separator_color)(x)
+    return offset_separator_color.color_function(x)
 
 
 def delimiter(x: object) -> str:
-    return generateColorFunction(config.telescope_offset_delimiter_color)(x)
+    return offset_delimiter_color.color_function(x)
 
 
 def repeating_marker(x: object) -> str:
-    return generateColorFunction(config.telescope_repeating_marker_color)(x)
+    return repeating_marker_color.color_function(x)

pwndbg/color/theme.py

@@ -1,15 +1,24 @@
 from __future__ import annotations
 
 from typing import Any
+from typing import Callable
 from typing import Sequence
 
+import pwndbg.color
 from pwndbg import config
 from pwndbg.lib.config import Parameter
 from pwndbg.lib.config import Scope
 
 
 class ColorParameter(Parameter):
-    pass
+    color_function: Callable[[object], str]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.update_color_function()
+
+    def update_color_function(self):
+        self.color_function = pwndbg.color.generateColorFunction(self.value)
 
 
 def add_param(
@@ -32,5 +41,11 @@ def add_param(
     )
 
 
-def add_color_param(name: str, default: Any, set_show_doc: str) -> Parameter:
-    return config.add_param_obj(ColorParameter(name, default, set_show_doc, scope=Scope.theme))
+def add_color_param(name: str, default: Any, set_show_doc: str) -> ColorParameter:
+    color_parameter = ColorParameter(name, default, set_show_doc, scope=Scope.theme)
+
+    config.triggers[name].append(color_parameter.update_color_function)
+
+    config.add_param_obj(color_parameter)
+
+    return color_parameter