From 87aafd15e727ba57b41d5285b5d3f485ac0ad180 Mon Sep 17 00:00:00 2001 From: OBarronCS <55004530+OBarronCS@users.noreply.github.com> Date: Fri, 16 May 2025 03:48:23 -0700 Subject: [PATCH] Cache string coloring functions (#2977) * Cache color functions until the color parameter is set * Update pwndbg/color/theme.py Co-authored-by: Disconnect3d --------- Co-authored-by: Disconnect3d --- pwndbg/color/__init__.py | 4 ++-- pwndbg/color/context.py | 26 ++++++++++++-------------- pwndbg/color/enhance.py | 10 ++++------ pwndbg/color/hexdump.py | 14 +++++++------- pwndbg/color/message.py | 39 +++++++++++++++++---------------------- pwndbg/color/telescope.py | 12 +++++------- pwndbg/color/theme.py | 21 ++++++++++++++++++--- 7 files changed, 65 insertions(+), 61 deletions(-) diff --git a/pwndbg/color/__init__.py b/pwndbg/color/__init__.py index 94d85121c..2bc4f9cdf 100644 --- a/pwndbg/color/__init__.py +++ b/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}'") diff --git a/pwndbg/color/context.py b/pwndbg/color/context.py index 36a12c415..cdd16cc68 100644 --- a/pwndbg/color/context.py +++ b/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): diff --git a/pwndbg/color/enhance.py b/pwndbg/color/enhance.py index 100664e54..ff68d46b7 100644 --- a/pwndbg/color/enhance.py +++ b/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) diff --git a/pwndbg/color/hexdump.py b/pwndbg/color/hexdump.py index 3930e80b2..5c5a98039 100644 --- a/pwndbg/color/hexdump.py +++ b/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: diff --git a/pwndbg/color/message.py b/pwndbg/color/message.py index 022c08a78..509b384f2 100644 --- a/pwndbg/color/message.py +++ b/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: diff --git a/pwndbg/color/telescope.py b/pwndbg/color/telescope.py index 84a39ab23..ae579614b 100644 --- a/pwndbg/color/telescope.py +++ b/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) diff --git a/pwndbg/color/theme.py b/pwndbg/color/theme.py index 7182922ca..f812c9571 100644 --- a/pwndbg/color/theme.py +++ b/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