diff --git a/README.rst b/README.rst index 0628b960..7d6f1b9a 100644 --- a/README.rst +++ b/README.rst @@ -84,3 +84,11 @@ version control tool.:: git clone https://github.com/inducer/pudb.git You may also `browse the code `_ online. + + +Customize and Extend PuDB +------------------------- + +You can contribute your custom stringifiers, themes and shells under +`pudb/contrib` folder. Currently the process is streamlined for stringifiers, +while shells and themes will require some refactoring of the core PuDB code. diff --git a/pudb/contrib/README.md b/pudb/contrib/README.md new file mode 100644 index 00000000..fed070bc --- /dev/null +++ b/pudb/contrib/README.md @@ -0,0 +1,13 @@ +# Community contributed extensions for pudb + +Here the community can extend pudb with custom stringifiers, themes and shells. + + +## How to contribute your stringifiers + +Simply add a new python module inside `contrib/stringifiers` that contains your custom stringifier. + +Then add your stringifier to the `CONTRIB_STRINGIFIERS` dict inside +`contrib/stringifiers/__init__.py`. + +The new options should appear in the pudb settings pane after setting the `Enable community contributed content` option. diff --git a/pudb/contrib/__init__.py b/pudb/contrib/__init__.py new file mode 100644 index 00000000..5d3b3760 --- /dev/null +++ b/pudb/contrib/__init__.py @@ -0,0 +1 @@ +from pudb.contrib.stringifiers import CONTRIB_STRINGIFIERS diff --git a/pudb/contrib/stringifiers/__init__.py b/pudb/contrib/stringifiers/__init__.py new file mode 100644 index 00000000..88e595f4 --- /dev/null +++ b/pudb/contrib/stringifiers/__init__.py @@ -0,0 +1,8 @@ +from pudb.contrib.stringifiers.torch_stringifier import torch_stringifier_fn + +CONTRIB_STRINGIFIERS = { + # User contributed stringifiers + # Use the contrib prefix for all keys to avoid clashes with the core stringifiers + # and make known to the user that this is community contributed code + "contrib.pytorch": torch_stringifier_fn, +} diff --git a/pudb/contrib/stringifiers/torch_stringifier.py b/pudb/contrib/stringifiers/torch_stringifier.py new file mode 100644 index 00000000..9189af47 --- /dev/null +++ b/pudb/contrib/stringifiers/torch_stringifier.py @@ -0,0 +1,38 @@ +from typing import Any + +try: + import torch + + HAVE_TORCH = 1 +except: + HAVE_TORCH = 0 + + +from pudb.var_view import default_stringifier + + +def torch_stringifier_fn(value: Any) -> str: + if not HAVE_TORCH: + # Fall back to default stringifier + + return default_stringifier(value) + + if isinstance(value, torch.nn.Module): + device: str = str(next(value.parameters()).device) + params: int = sum([p.numel() for p in value.parameters() if p.requires_grad]) + rep: str = value.__repr__() if len(value.__repr__()) < 55 else type( + value + ).__name__ + + return "{}[{}] Params: {}".format(rep, device, params) + elif isinstance(value, torch.Tensor): + return "{}[{}][{}] {}".format( + type(value).__name__, + str(value.dtype).replace("torch.", ""), + str(value.device), + str(list(value.shape)), + ) + else: + # Fall back to default stringifier + + return default_stringifier(value) diff --git a/pudb/settings.py b/pudb/settings.py index 5a6d9a38..b95e91f5 100644 --- a/pudb/settings.py +++ b/pudb/settings.py @@ -25,7 +25,6 @@ import os import sys - from configparser import ConfigParser from pudb.lowlevel import (lookup_module, get_breakpoint_invalid_reason, settings_log) @@ -123,6 +122,8 @@ def load_config(): conf_dict.setdefault("hide_cmdline_win", "False") + conf_dict.setdefault("enable_community_contributed_content", "False") + def normalize_bool_inplace(name): try: if conf_dict[name].lower() in ["0", "false", "off"]: @@ -136,6 +137,7 @@ def normalize_bool_inplace(name): normalize_bool_inplace("wrap_variables") normalize_bool_inplace("prompt_on_quit") normalize_bool_inplace("hide_cmdline_win") + normalize_bool_inplace("enable_community_contributed_content") _config_[0] = conf_dict return conf_dict @@ -223,6 +225,11 @@ def _update_config(check_box, new_state, option_newvalue): conf_dict.update(new_conf_dict) _update_hide_cmdline_win() + elif option == "enable_community_contributed_content": + new_conf_dict["enable_community_contributed_content"] = ( + not check_box.get_state()) + conf_dict.update(new_conf_dict) + elif option == "current_stack_frame": # only activate if the new state of the radio button is 'on' if new_state: @@ -270,6 +277,14 @@ def _update_config(check_box, new_state, option_newvalue): bool(conf_dict["hide_cmdline_win"]), on_state_change=_update_config, user_data=("hide_cmdline_win", None)) + enable_community_contributed_content = urwid.CheckBox( + "Enable community contributed content. This will give you access to more " + "stringifiers, shells and themes. \n" + "Changing this setting requires a restart of PuDB.", + bool(conf_dict["enable_community_contributed_content"]), + on_state_change=_update_config, + user_data=("enable_community_contributed_content", None)) + # {{{ shells shell_info = urwid.Text("This is the shell that will be " @@ -345,8 +360,18 @@ def _update_config(check_box, new_state, option_newvalue): # {{{ stringifier from pudb.var_view import STRINGIFIERS + from pudb.contrib.stringifiers import CONTRIB_STRINGIFIERS stringifier_opts = list(STRINGIFIERS.keys()) + if conf_dict["enable_community_contributed_content"]: + stringifier_opts = ( + list(STRINGIFIERS.keys()) + list(CONTRIB_STRINGIFIERS.keys())) known_stringifier = conf_dict["stringifier"] in stringifier_opts + contrib_stringifier = conf_dict["stringifier"] in CONTRIB_STRINGIFIERS + fallback_to_default_stringifier = (contrib_stringifier and not + conf_dict["enable_community_contributed_content"]) + use_default_stringifier = ((conf_dict["stringifier"] == "default") or + fallback_to_default_stringifier) + custom_stringifier = not (known_stringifier or contrib_stringifier) stringifier_rb_group = [] stringifier_edit = urwid.Edit(edit_text=conf_dict["custom_stringifier"]) stringifier_info = urwid.Text( @@ -357,15 +382,21 @@ def _update_config(check_box, new_state, option_newvalue): "be slower than the default, type, or id stringifiers.\n") stringifier_edit_list_item = urwid.AttrMap(stringifier_edit, "input", "focused input") + stringifier_rbs = [ + urwid.RadioButton(stringifier_rb_group, "default", + use_default_stringifier, + on_state_change=_update_config, + user_data=("stringifier", "default")) + ]+[ urwid.RadioButton(stringifier_rb_group, name, conf_dict["stringifier"] == name, on_state_change=_update_config, user_data=("stringifier", name)) - for name in stringifier_opts + for name in stringifier_opts if name != "default" ]+[ urwid.RadioButton(stringifier_rb_group, "Custom:", - not known_stringifier, on_state_change=_update_config, + custom_stringifier, on_state_change=_update_config, user_data=("stringifier", None)), stringifier_edit_list_item, urwid.Text("\nTo use a custom stringifier, see " @@ -441,6 +472,7 @@ def _update_config(check_box, new_state, option_newvalue): + [cb_line_numbers] + [cb_prompt_on_quit] + [hide_cmdline_win] + + [enable_community_contributed_content] + [urwid.AttrMap(urwid.Text("\nShell:\n"), "group head")] + [shell_info] diff --git a/pudb/var_view.py b/pudb/var_view.py index 843cab44..bce3d1d1 100644 --- a/pudb/var_view.py +++ b/pudb/var_view.py @@ -203,7 +203,8 @@ def __init__(self): self.access_level = CONFIG["default_variables_access_level"] self.show_methods = False self.wrap = CONFIG["wrap_variables"] - + self.enable_contrib_stringifiers = \ + CONFIG["enable_community_contributed_content"] class WatchExpression: def __init__(self, expression): @@ -456,11 +457,19 @@ def error_stringifier(_): } +from pudb.contrib.stringifiers import CONTRIB_STRINGIFIERS + + def get_stringifier(iinfo: InspectInfo) -> Callable: """ :return: a function that turns an object into a Unicode text object. """ try: + if iinfo.display_type in CONTRIB_STRINGIFIERS: + if iinfo.enable_contrib_stringifiers: + return CONTRIB_STRINGIFIERS[iinfo.display_type] + else: + return STRINGIFIERS["default"] return STRINGIFIERS[iinfo.display_type] except KeyError: try: diff --git a/test/test_contrib_torch_stringifier.py b/test/test_contrib_torch_stringifier.py new file mode 100644 index 00000000..8f40d246 --- /dev/null +++ b/test/test_contrib_torch_stringifier.py @@ -0,0 +1,46 @@ +try: + import torch + HAVE_TORCH = True +except ImportError: + HAVE_TORCH = False + +from pudb.var_view import default_stringifier +from pudb.contrib.stringifiers.torch_stringifier import torch_stringifier_fn + +def test_tensor(): + if HAVE_TORCH: + x = torch.randn(10, 5, 4) + assert torch_stringifier_fn(x) == "Tensor[float32][cpu] [10, 5, 4]" + + +def test_conv_module(): + if HAVE_TORCH: + x = torch.nn.Conv2d(20, 10, 3) + assert torch_stringifier_fn(x) == "Conv2d(20, 10, kernel_size=(3, 3), stride=(1, 1))[cpu] Params: 1810" + + +def test_linear_module(): + if HAVE_TORCH: + x = torch.nn.Linear(5, 2, bias=False) + assert torch_stringifier_fn(x) == "Linear(in_features=5, out_features=2, bias=False)[cpu] Params: 10" + + +def test_long_module_repr_should_revert_to_type(): + if HAVE_TORCH: + x = torch.nn.Transformer() + assert torch_stringifier_fn(x) == "Transformer[cpu] Params: 44140544" + + +def test_reverts_to_default_for_str(): + x = "Everyone has his day, and some days last longer than others." + assert torch_stringifier_fn(x) == default_stringifier(x) + + +def test_reverts_to_default_for_dict(): + x = {"a": 1, "b": 2, "c": 3} + assert torch_stringifier_fn(x) == default_stringifier(x) + + +def test_reverts_to_default_for_list(): + x = list(range(1000)) + assert torch_stringifier_fn(x) == default_stringifier(x)