tfdbg: add tab completion to ReadlineUI

Change: 144108181

3 files changed, 59 insertions, 24 deletions

diff --git a/tensorflow/python/debug/cli/base_ui.py b/tensorflow/python/debug/cli/base_ui.py
index 30e43e81ff..eed2b03a1a 100644
--- a/tensorflow/python/debug/cli/base_ui.py
+++ b/tensorflow/python/debug/cli/base_ui.py
@@ -134,3 +134,45 @@ class BaseUI(object):
         command_items)
 
     return command_items[0], command_items[1:], output_file_path
+
+  def _analyze_tab_complete_input(self, text):
+    """Analyze raw input to tab-completer.
+
+    Args:
+      text: (str) the full, raw input text to be tab-completed.
+
+    Returns:
+      context: (str) the context str. For example,
+        If text == "print_tensor softmax", returns "print_tensor".
+        If text == "print", returns "".
+        If text == "", returns "".
+      prefix: (str) the prefix to be tab-completed, from the last word.
+        For example, if text == "print_tensor softmax", returns "softmax".
+        If text == "print", returns "print".
+        If text == "", returns "".
+      except_last_word: (str) the input text, except the last word.
+        For example, if text == "print_tensor softmax", returns "print_tensor".
+        If text == "print_tensor -a softmax", returns "print_tensor -a".
+        If text == "print", returns "".
+        If text == "", returns "".
+    """
+    text = text.lstrip()
+    if not text:
+      # Empty (top-level) context.
+      context = ""
+      prefix = ""
+      except_last_word = ""
+    else:
+      items = text.split(" ")
+      if len(items) == 1:
+        # Single word: top-level context.
+        context = ""
+        prefix = items[0]
+        except_last_word = ""
+      else:
+        # Multiple words.
+        context = items[0]
+        prefix = items[-1]
+        except_last_word = " ".join(items[:-1]) + " "
+
+    return context, prefix, except_last_word
diff --git a/tensorflow/python/debug/cli/curses_ui.py b/tensorflow/python/debug/cli/curses_ui.py
index fa7ec280fe..ba17c1632d 100644
--- a/tensorflow/python/debug/cli/curses_ui.py
+++ b/tensorflow/python/debug/cli/curses_ui.py
@@ -1154,24 +1154,8 @@ class CursesUI(base_ui.BaseUI):
       appended by the common prefix of the candidates.
     """
 
-    command_str = command_str.lstrip()
-
-    if not command_str:
-      # Empty (top-level) context.
-      context = ""
-      prefix = ""
-      items = []
-    else:
-      items = command_str.split(" ")
-      if len(items) == 1:
-        # Single word: top-level context.
-        context = ""
-        prefix = items[0]
-      else:
-        # Multiple words.
-        context = items[0]
-        prefix = items[-1]
-
+    context, prefix, except_last_word = self._analyze_tab_complete_input(
+        command_str)
     candidates, common_prefix = self._tab_completion_registry.get_completions(
         context, prefix)
 
@@ -1186,9 +1170,9 @@ class CursesUI(base_ui.BaseUI):
     if common_prefix:
       # Common prefix is not None and non-empty. The completed string will
       # incorporate the common prefix.
-      return " ".join(items[:-1] + [common_prefix])
+      return except_last_word + common_prefix
     else:
-      return " ".join(items)
+      return except_last_word + prefix
 
   def _display_candidates(self, candidates):
     """Show candidates (e.g., tab-completion candidates) on multiple lines.
diff --git a/tensorflow/python/debug/cli/readline_ui.py b/tensorflow/python/debug/cli/readline_ui.py
index 850ebf36f9..4317d14461 100644
--- a/tensorflow/python/debug/cli/readline_ui.py
+++ b/tensorflow/python/debug/cli/readline_ui.py
@@ -28,21 +28,30 @@ class ReadlineUI(base_ui.BaseUI):
 
   def __init__(self, on_ui_exit=None):
     base_ui.BaseUI.__init__(self, on_ui_exit=on_ui_exit)
-
     self._init_input()
-    # TODO(cais): Enable readline tab completion support via
-    #   self._tab_completion_registry.
 
   def _init_input(self):
-    readline.parse_and_bind("tab: complete")
     readline.parse_and_bind("set editing-mode emacs")
 
+    # Disable default readline delimiter in order to receive the full text
+    # (not just the last word) in the completer.
+    readline.set_completer_delims("\n")
+    readline.set_completer(self._readline_complete)
+    readline.parse_and_bind("tab: complete")
+
     # For Python 2-3 compatibility.
     try:
       self._input = raw_input
     except NameError:
       self._input = input
 
+  def _readline_complete(self, text, state):
+    context, prefix, except_last_word = self._analyze_tab_complete_input(text)
+    candidates, _ = self._tab_completion_registry.get_completions(context,
+                                                                  prefix)
+    candidates = [(except_last_word + candidate) for candidate in candidates]
+    return candidates[state]
+
   def run_ui(self,
              init_command=None,
              title=None,