Include contrib folder and pytorch stringifier

This PR introduces a new contrib folder under pudb for community
customization. A stringifier for pytorch is included.

Detailed changelog:

* A module was added under pudb/contrib/stringifiers as place to store
  custom stringifiers.
* A stringifier for pytorch tensors and modules is included.
* An option was added in the settings menu for the user to enable or
  disable the cotrib content for users that want to stick to the core
  pudb installation.
* Changes were made in var_view.py and settings.py to allow for
  inclusion of the contrib/stringifiers in the configuration menu.

Signed-off-by: Giorgos Paraskevopoulos <[email protected]>

Author: georgepar

README.rst

@@ -84,3 +84,11 @@ version control tool.::
     git clone https://github.com/inducer/pudb.git
 
 You may also `browse the code <https://github.com/inducer/pudb>`_ 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.

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.

pudb/contrib/__init__.py

@@ -0,0 +1 @@
+from pudb.contrib.stringifiers import CONTRIB_STRINGIFIERS

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,
+}

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)

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]

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:

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)