xref: /aosp_15_r20/external/tensorflow/tensorflow/python/debug/cli/base_ui.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Base Class of TensorFlow Debugger (tfdbg) Command-Line Interface."""
16import argparse
17
18from tensorflow.python.debug.cli import cli_config
19from tensorflow.python.debug.cli import command_parser
20from tensorflow.python.debug.cli import debugger_cli_common
21
22
23class BaseUI(object):
24  """Base class of tfdbg user interface."""
25
26  CLI_PROMPT = "tfdbg> "
27  CLI_EXIT_COMMANDS = ["exit", "quit"]
28  ERROR_MESSAGE_PREFIX = "ERROR: "
29  INFO_MESSAGE_PREFIX = "INFO: "
30
31  def __init__(self, on_ui_exit=None, config=None):
32    """Constructor of the base class.
33
34    Args:
35      on_ui_exit: (`Callable`) the callback to be called when the UI exits.
36      config: An instance of `cli_config.CLIConfig()` carrying user-facing
37        configurations.
38    """
39
40    self._on_ui_exit = on_ui_exit
41
42    self._command_handler_registry = (
43        debugger_cli_common.CommandHandlerRegistry())
44
45    self._tab_completion_registry = debugger_cli_common.TabCompletionRegistry()
46
47    # Create top-level tab-completion context and register the exit and help
48    # commands.
49    self._tab_completion_registry.register_tab_comp_context(
50        [""], self.CLI_EXIT_COMMANDS +
51        [debugger_cli_common.CommandHandlerRegistry.HELP_COMMAND] +
52        debugger_cli_common.CommandHandlerRegistry.HELP_COMMAND_ALIASES)
53
54    self._config = config or cli_config.CLIConfig()
55    self._config_argparser = argparse.ArgumentParser(
56        description="config command", usage=argparse.SUPPRESS)
57    subparsers = self._config_argparser.add_subparsers()
58    set_parser = subparsers.add_parser("set")
59    set_parser.add_argument("property_name", type=str)
60    set_parser.add_argument("property_value", type=str)
61    set_parser = subparsers.add_parser("show")
62    self.register_command_handler(
63        "config",
64        self._config_command_handler,
65        self._config_argparser.format_help(),
66        prefix_aliases=["cfg"])
67
68  def set_help_intro(self, help_intro):
69    """Set an introductory message to the help output of the command registry.
70
71    Args:
72      help_intro: (RichTextLines) Rich text lines appended to the beginning of
73        the output of the command "help", as introductory information.
74    """
75
76    self._command_handler_registry.set_help_intro(help_intro=help_intro)
77
78  def register_command_handler(self,
79                               prefix,
80                               handler,
81                               help_info,
82                               prefix_aliases=None):
83    """A wrapper around CommandHandlerRegistry.register_command_handler().
84
85    In addition to calling the wrapped register_command_handler() method, this
86    method also registers the top-level tab-completion context based on the
87    command prefixes and their aliases.
88
89    See the doc string of the wrapped method for more details on the args.
90
91    Args:
92      prefix: (str) command prefix.
93      handler: (callable) command handler.
94      help_info: (str) help information.
95      prefix_aliases: (list of str) aliases of the command prefix.
96    """
97
98    self._command_handler_registry.register_command_handler(
99        prefix, handler, help_info, prefix_aliases=prefix_aliases)
100
101    self._tab_completion_registry.extend_comp_items("", [prefix])
102    if prefix_aliases:
103      self._tab_completion_registry.extend_comp_items("", prefix_aliases)
104
105  def register_tab_comp_context(self, *args, **kwargs):
106    """Wrapper around TabCompletionRegistry.register_tab_comp_context()."""
107
108    self._tab_completion_registry.register_tab_comp_context(*args, **kwargs)
109
110  def run_ui(self,
111             init_command=None,
112             title=None,
113             title_color=None,
114             enable_mouse_on_start=True):
115    """Run the UI until user- or command- triggered exit.
116
117    Args:
118      init_command: (str) Optional command to run on CLI start up.
119      title: (str) Optional title to display in the CLI.
120      title_color: (str) Optional color of the title, e.g., "yellow".
121      enable_mouse_on_start: (bool) Whether the mouse mode is to be enabled on
122        start-up.
123
124    Returns:
125      An exit token of arbitrary type. Can be None.
126    """
127
128    raise NotImplementedError("run_ui() is not implemented in BaseUI")
129
130  def _parse_command(self, command):
131    """Parse a command string into prefix and arguments.
132
133    Args:
134      command: (str) Command string to be parsed.
135
136    Returns:
137      prefix: (str) The command prefix.
138      args: (list of str) The command arguments (i.e., not including the
139        prefix).
140      output_file_path: (str or None) The path to save the screen output
141        to (if any).
142    """
143    command = command.strip()
144    if not command:
145      return "", [], None
146
147    command_items = command_parser.parse_command(command)
148    command_items, output_file_path = command_parser.extract_output_file_path(
149        command_items)
150
151    return command_items[0], command_items[1:], output_file_path
152
153  def _analyze_tab_complete_input(self, text):
154    """Analyze raw input to tab-completer.
155
156    Args:
157      text: (str) the full, raw input text to be tab-completed.
158
159    Returns:
160      context: (str) the context str. For example,
161        If text == "print_tensor softmax", returns "print_tensor".
162        If text == "print", returns "".
163        If text == "", returns "".
164      prefix: (str) the prefix to be tab-completed, from the last word.
165        For example, if text == "print_tensor softmax", returns "softmax".
166        If text == "print", returns "print".
167        If text == "", returns "".
168      except_last_word: (str) the input text, except the last word.
169        For example, if text == "print_tensor softmax", returns "print_tensor".
170        If text == "print_tensor -a softmax", returns "print_tensor -a".
171        If text == "print", returns "".
172        If text == "", returns "".
173    """
174    text = text.lstrip()
175    if not text:
176      # Empty (top-level) context.
177      context = ""
178      prefix = ""
179      except_last_word = ""
180    else:
181      items = text.split(" ")
182      if len(items) == 1:
183        # Single word: top-level context.
184        context = ""
185        prefix = items[0]
186        except_last_word = ""
187      else:
188        # Multiple words.
189        context = items[0]
190        prefix = items[-1]
191        except_last_word = " ".join(items[:-1]) + " "
192
193    return context, prefix, except_last_word
194
195  @property
196  def config(self):
197    """Obtain the CLIConfig of this `BaseUI` instance."""
198    return self._config
199
200  def _config_command_handler(self, args, screen_info=None):
201    """Command handler for the "config" command."""
202    del screen_info  # Currently unused.
203
204    parsed = self._config_argparser.parse_args(args)
205    if hasattr(parsed, "property_name") and hasattr(parsed, "property_value"):
206      # set.
207      self._config.set(parsed.property_name, parsed.property_value)
208      return self._config.summarize(highlight=parsed.property_name)
209    else:
210      # show.
211      return self._config.summarize()
212