move local var printing out of tool into locals.py

src/chatdbg/chatdbg_pdb.py

@@ -17,12 +17,12 @@
 
 from .assistant.assistant import Assistant
 from .pdb.capture import CaptureInput, CaptureOutput
+from .pdb.locals import print_locals
+from .pdb.prompts import pdb_instructions
+from .pdb.text import strip_color, truncate_proportionally
 from .util.config import chatdbg_config
-from .pdb.locals import extract_locals
 from .util.log import ChatDBGLog
 from .util.printer import ChatDBGPrinter
-from .pdb.prompts import pdb_instructions
-from .pdb.text import (format_limited, strip_color, truncate_proportionally)
 
 
 def load_ipython_extension(ipython):
@@ -131,7 +131,7 @@ def _is_user_file(self, file_name):
 
         return False
 
-    def format_stack_trace(self, context=None):
+    def enriched_stack_trace(self, context=None):
         old_stdout = self.stdout
         buf = StringIO()
         self.stdout = buf
@@ -188,7 +188,7 @@ def execRcLines(self):
         # do before running rclines -- our stack should be set up by now.
         if not chatdbg_config.show_libs:
             self._hide_lib_frames()
-        self._error_stack_trace = f"The program has the following stack trace:\n```\n{self.format_stack_trace()}\n```\n"
+        self._error_stack_trace = f"The program has the following stack trace:\n```\n{self.enriched_stack_trace()}\n```\n"
 
         # finally safe to enable this.
         self._show_locals = chatdbg_config.show_locals and not chatdbg_config.show_libs
@@ -462,7 +462,7 @@ def print_stack_trace(self, context=None, locals=None):
                     skipped = 0
                 self.print_stack_entry(frame_lineno, context=context)
                 if locals:
-                    self._print_locals(frame_lineno[0])
+                    print_locals(self.stdout, frame_lineno[0])
             if skipped:
                 print(
                     f"{Colors.excName}    [... skipping {skipped} hidden frame(s)]{ColorsNormal}\n",
@@ -471,34 +471,6 @@ def print_stack_trace(self, context=None, locals=None):
         except KeyboardInterrupt:
             pass
 
-    def _print_locals(self, frame):
-        locals = frame.f_locals
-        in_global_scope = locals is frame.f_globals
-        defined_locals = extract_locals(frame)
-        # Unclear benefit: possibly some benefit w/ stack only runs, but large context...
-        # if in_global_scope and "In" in locals:  # in notebook
-        #     defined_locals = defined_locals | extract_nb_globals(locals)
-        if len(defined_locals) > 0:
-            if in_global_scope:
-                print(f"    Global variables:", file=self.stdout)
-            else:
-                print(f"    Variables in this frame:", file=self.stdout)
-            for name in sorted(defined_locals):
-                value = locals[name]
-                t = type(value).__name__
-                prefix = f"      {name}: {t} = "
-                rep = format_limited(value, limit=20).split("\n")
-                if len(rep) > 1:
-                    rep = (
-                        prefix
-                        + rep[0]
-                        + "\n"
-                        + textwrap.indent("\n".join(rep[1:]), prefix=" " * len(prefix))
-                    )
-                else:
-                    rep = prefix + rep[0]
-                print(rep, file=self.stdout)
-            print(file=self.stdout)
 
     def _stack_prompt(self):
         stdout = self.stdout
@@ -523,16 +495,16 @@ def _stack_prompt(self):
         finally:
             self.stdout = stdout
 
-    def _ip_instructions(self):
+    def _initial_prompt_instructions(self):
         return pdb_instructions(self._supports_flow, chatdbg_config.take_the_wheel)
 
-    def _ip_enchriched_stack_trace(self):
-        return f"The program has this stack trace:\n```\n{self.format_stack_trace()}\n```\n"
+    def _initial_prompt_enchriched_stack_trace(self):
+        return f"The program has this stack trace:\n```\n{self.enriched_stack_trace()}\n```\n"
 
-    def _ip_error(self):
+    def _initial_prompt_error(self):
         return self._error_specific_prompt
 
-    def _ip_inputs(self):
+    def _initial_prompt_inputs(self):
         inputs = ""
         if len(sys.argv) > 1:
             inputs += f"\nThese were the command line options:\n```\n{' '.join(sys.argv)}\n```\n"
@@ -541,7 +513,7 @@ def _ip_inputs(self):
             inputs += f"\nThis was the program's input :\n```\n{input}```\n"
         return inputs
 
-    def _ip_history(self):
+    def _initial_prompt_history(self):
         if len(self._history) > 0:
             hist = textwrap.indent(self._capture_onecmd("hist"), "")
             hist = f"\nThis is the history of some pdb commands I ran and the results.\n```\n{hist}\n```\n"
@@ -556,14 +528,14 @@ def concat_prompt(self, *args):
     def _build_prompt(self, arg, conversing):
         if not conversing:
             return self.concat_prompt(
-                self._ip_enchriched_stack_trace(),
-                self._ip_inputs(),
-                self._ip_error(),
-                self._ip_history(),
+                self._initial_prompt_enchriched_stack_trace(),
+                self._initial_prompt_inputs(),
+                self._initial_prompt_error(),
+                self._initial_prompt_history(),
                 arg,
             )
         else:
-            return self.concat_prompt(self._ip_history(), self._stack_prompt(), arg)
+            return self.concat_prompt(self._initial_prompt_history(), self._stack_prompt(), arg)
 
     def do_chat(self, arg):
         """chat
@@ -595,6 +567,10 @@ def do_renew(self, arg):
         self.message(f"Ready to start a new dialog.")
 
     def do_config(self, arg):
+        """
+        config
+        Print out the ChatDBG config options.
+        """
         args = arg.split()
         if len(args) == 0:
             pprint(chatdbg_config.to_json(), sort_dicts=True, stream=self.stdout)
@@ -613,7 +589,7 @@ def do_config(self, arg):
             self.error(f"{e}")
 
     def _make_assistant(self):
-        instruction_prompt = self._ip_instructions()
+        instruction_prompt = self._initial_prompt_instructions()
 
         if chatdbg_config.take_the_wheel:
             functions = [self.debug, self.info]

src/chatdbg/pdb/locals.py

@@ -1,7 +1,11 @@
 import ast
 import inspect
+import itertools
+import numbers
 import textwrap
 
+import numpy as np
+
 
 class SymbolFinder(ast.NodeVisitor):
     def __init__(self):
@@ -24,7 +28,7 @@ def visit_comprehension(self, node):
         self.generic_visit(node)
 
 
-def extract_locals(frame):
+def _extract_locals(frame):
     try:
         source = textwrap.dedent(inspect.getsource(frame))
         tree = ast.parse(source)
@@ -42,7 +46,7 @@ def extract_locals(frame):
         return set()
 
 
-def extract_nb_globals(globals):
+def _extract_nb_globals(globals):
     result = set()
     for source in globals["In"]:
         try:
@@ -53,3 +57,123 @@ def extract_nb_globals(globals):
         except Exception as e:
             pass
     return result
+
+
+
+def _is_iterable(obj):
+    try:
+        iter(obj)
+        return True
+    except TypeError:
+        return False
+
+
+def _repr_if_defined(obj):
+    if obj.__class__ in [np.ndarray, dict, list, tuple]:
+        # handle these at iterables to truncate reasonably
+        return False
+    result = (
+        "__repr__" in dir(obj.__class__)
+        and obj.__class__.__repr__ is not object.__repr__
+    )
+    return result
+
+
+def _format_limited(value, limit=10, depth=3):
+    def format_tuple(t, depth):
+        return tuple([helper(x, depth) for x in t])
+
+    def format_list(list, depth):
+        return [helper(x, depth) for x in list]
+
+    def format_dict(items, depth):
+        return {k: helper(v, depth) for k, v in items}
+
+    def format_object(obj, depth):
+        attributes = dir(obj)
+        fields = {
+            attr: getattr(obj, attr, None)
+            for attr in attributes
+            if not callable(getattr(obj, attr, None)) and not attr.startswith("__")
+        }
+        return format(
+            f"{type(obj).__name__} object with fields {format_dict(fields.items(), depth)}"
+        )
+
+    def helper(value, depth):
+        if depth == 0:
+            return ...
+        if value is Ellipsis:
+            return ...
+        if isinstance(value, dict):
+            if len(value) > limit:
+                return format_dict(
+                    list(value.items())[: limit - 1] + [(..., ...)], depth - 1
+                )
+            else:
+                return format_dict(value.items(), depth - 1)
+        elif isinstance(value, (str, bytes)):
+            if len(value) > 254:
+                value = str(value)[0:253] + "..."
+            return value
+        elif isinstance(value, tuple):
+            if len(value) > limit:
+                return format_tuple(value[0 : limit - 1] + (...,), depth - 1)
+            else:
+                return format_tuple(value, depth - 1)
+        elif value is None or isinstance(
+            value, (int, float, bool, type, numbers.Number)
+        ):
+            return value
+        elif isinstance(value, np.ndarray):
+            with np.printoptions(threshold=limit):
+                return np.array_repr(value)
+        elif inspect.isclass(type(value)) and _repr_if_defined(value):
+            return repr(value)
+        elif _is_iterable(value):
+            value = list(itertools.islice(value, 0, limit + 1))
+            if len(value) > limit:
+                return format_list(value[: limit - 1] + [...], depth - 1)
+            else:
+                return format_list(value, depth - 1)
+        elif inspect.isclass(type(value)):
+            return format_object(value, depth - 1)
+        else:
+            return value
+
+    result = str(helper(value, depth=3)).replace("Ellipsis", "...")
+    if len(result) > 1024 * 2:
+        result = result[: 1024 * 2 - 3] + "..."
+    if type(value) == str:
+        return "'" + result + "'"
+    else:
+        return result
+
+def print_locals(file, frame):
+    locals = frame.f_locals
+    in_global_scope = locals is frame.f_globals
+    defined_locals = _extract_locals(frame)
+    # Unclear benefit: possibly some benefit w/ stack only runs, but large context...
+    # if in_global_scope and "In" in locals:  # in notebook
+    #     defined_locals = defined_locals | extract_nb_globals(locals)
+    if len(defined_locals) > 0:
+        if in_global_scope:
+            print(f"    Global variables:", file=file)
+        else:
+            print(f"    Variables in this frame:", file=file)
+        for name in sorted(defined_locals):
+            value = locals[name]
+            t = type(value).__name__
+            prefix = f"      {name}: {t} = "
+            rep = _format_limited(value, limit=20).split("\n")
+            if len(rep) > 1:
+                rep = (
+                    prefix
+                    + rep[0]
+                    + "\n"
+                    + textwrap.indent("\n".join(rep[1:]), prefix=" " * len(prefix))
+                )
+            else:
+                rep = prefix + rep[0]
+            print(rep, file=file)
+        print(file=file)