xref: /aosp_15_r20/external/tensorflow/tensorflow/python/debug/cli/debugger_cli_common.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"""Building Blocks of TensorFlow Debugger Command-Line Interface."""
16import copy
17import os
18import re
19import sre_constants
20import traceback
21
22import numpy as np
23
24from tensorflow.python.client import pywrap_tf_session
25from tensorflow.python.platform import gfile
26
27HELP_INDENT = "  "
28
29EXPLICIT_USER_EXIT = "explicit_user_exit"
30REGEX_MATCH_LINES_KEY = "regex_match_lines"
31INIT_SCROLL_POS_KEY = "init_scroll_pos"
32
33MAIN_MENU_KEY = "mm:"
34
35
36class CommandLineExit(Exception):
37
38  def __init__(self, exit_token=None):
39    Exception.__init__(self)
40    self._exit_token = exit_token
41
42  @property
43  def exit_token(self):
44    return self._exit_token
45
46
47class RichLine:
48  """Rich single-line text.
49
50  Attributes:
51    text: A plain string, the raw text represented by this object.  Should not
52      contain newlines.
53    font_attr_segs: A list of (start, end, font attribute) triples, representing
54      richness information applied to substrings of text.
55  """
56
57  def __init__(self, text="", font_attr=None):
58    """Construct a RichLine with no rich attributes or a single attribute.
59
60    Args:
61      text: Raw text string
62      font_attr: If specified, a single font attribute to be applied to the
63        entire text.  Extending this object via concatenation allows creation
64        of text with varying attributes.
65    """
66    # TODO(ebreck) Make .text and .font_attr protected members when we no
67    # longer need public access.
68    self.text = text
69    if font_attr:
70      self.font_attr_segs = [(0, len(text), font_attr)]
71    else:
72      self.font_attr_segs = []
73
74  def __add__(self, other):
75    """Concatenate two chunks of maybe rich text to make a longer rich line.
76
77    Does not modify self.
78
79    Args:
80      other: Another piece of text to concatenate with this one.
81        If it is a plain str, it will be appended to this string with no
82        attributes.  If it is a RichLine, it will be appended to this string
83        with its attributes preserved.
84
85    Returns:
86      A new RichLine comprising both chunks of text, with appropriate
87        attributes applied to the corresponding substrings.
88    """
89    ret = RichLine()
90    if isinstance(other, str):
91      ret.text = self.text + other
92      ret.font_attr_segs = self.font_attr_segs[:]
93      return ret
94    elif isinstance(other, RichLine):
95      ret.text = self.text + other.text
96      ret.font_attr_segs = self.font_attr_segs[:]
97      old_len = len(self.text)
98      for start, end, font_attr in other.font_attr_segs:
99        ret.font_attr_segs.append((old_len + start, old_len + end, font_attr))
100      return ret
101    else:
102      raise TypeError("%r cannot be concatenated with a RichLine" % other)
103
104  def __len__(self):
105    return len(self.text)
106
107
108def rich_text_lines_from_rich_line_list(rich_text_list, annotations=None):
109  """Convert a list of RichLine objects or strings to a RichTextLines object.
110
111  Args:
112    rich_text_list: a list of RichLine objects or strings
113    annotations: annotations for the resultant RichTextLines object.
114
115  Returns:
116    A corresponding RichTextLines object.
117  """
118  lines = []
119  font_attr_segs = {}
120  for i, rl in enumerate(rich_text_list):
121    if isinstance(rl, RichLine):
122      lines.append(rl.text)
123      if rl.font_attr_segs:
124        font_attr_segs[i] = rl.font_attr_segs
125    else:
126      lines.append(rl)
127  return RichTextLines(lines, font_attr_segs, annotations=annotations)
128
129
130def get_tensorflow_version_lines(include_dependency_versions=False):
131  """Generate RichTextLines with TensorFlow version info.
132
133  Args:
134    include_dependency_versions: Include the version of TensorFlow's key
135      dependencies, such as numpy.
136
137  Returns:
138    A formatted, multi-line `RichTextLines` object.
139  """
140  lines = ["TensorFlow version: %s" % pywrap_tf_session.__version__]
141  lines.append("")
142  if include_dependency_versions:
143    lines.append("Dependency version(s):")
144    lines.append("  numpy: %s" % np.__version__)
145    lines.append("")
146  return RichTextLines(lines)
147
148
149class RichTextLines:
150  """Rich multi-line text.
151
152  Line-by-line text output, with font attributes (e.g., color) and annotations
153  (e.g., indices in a multi-dimensional tensor). Used as the text output of CLI
154  commands. Can be rendered on terminal environments such as curses.
155
156  This is not to be confused with Rich Text Format (RTF). This class is for text
157  lines only.
158  """
159
160  def __init__(self, lines, font_attr_segs=None, annotations=None):
161    """Constructor of RichTextLines.
162
163    Args:
164      lines: A list of str or a single str, representing text output to
165        screen. The latter case is for convenience when the text output is
166        single-line.
167      font_attr_segs: A map from 0-based row index to a list of 3-tuples.
168        It lists segments in each row that have special font attributes, such
169        as colors, that are not the default attribute. For example:
170        {1: [(0, 3, "red"), (4, 7, "green")], 2: [(10, 20, "yellow")]}
171
172        In each tuple, the 1st element is the start index of the segment. The
173        2nd element is the end index, in an "open interval" fashion. The 3rd
174        element is an object or a list of objects that represents the font
175        attribute. Colors are represented as strings as in the examples above.
176      annotations: A map from 0-based row index to any object for annotating
177        the row. A typical use example is annotating rows of the output as
178        indices in a multi-dimensional tensor. For example, consider the
179        following text representation of a 3x2x2 tensor:
180          [[[0, 0], [0, 0]],
181           [[0, 0], [0, 0]],
182           [[0, 0], [0, 0]]]
183        The annotation can indicate the indices of the first element shown in
184        each row, i.e.,
185          {0: [0, 0, 0], 1: [1, 0, 0], 2: [2, 0, 0]}
186        This information can make display of tensors on screen clearer and can
187        help the user navigate (scroll) to the desired location in a large
188        tensor.
189
190    Raises:
191      ValueError: If lines is of invalid type.
192    """
193    if isinstance(lines, list):
194      self._lines = lines
195    elif isinstance(lines, str):
196      self._lines = [lines]
197    else:
198      raise ValueError("Unexpected type in lines: %s" % type(lines))
199
200    self._font_attr_segs = font_attr_segs
201    if not self._font_attr_segs:
202      self._font_attr_segs = {}
203      # TODO(cais): Refactor to collections.defaultdict(list) to simplify code.
204
205    self._annotations = annotations
206    if not self._annotations:
207      self._annotations = {}
208      # TODO(cais): Refactor to collections.defaultdict(list) to simplify code.
209
210  @property
211  def lines(self):
212    return self._lines
213
214  @property
215  def font_attr_segs(self):
216    return self._font_attr_segs
217
218  @property
219  def annotations(self):
220    return self._annotations
221
222  def num_lines(self):
223    return len(self._lines)
224
225  def slice(self, begin, end):
226    """Slice a RichTextLines object.
227
228    The object itself is not changed. A sliced instance is returned.
229
230    Args:
231      begin: (int) Beginning line index (inclusive). Must be >= 0.
232      end: (int) Ending line index (exclusive). Must be >= 0.
233
234    Returns:
235      (RichTextLines) Sliced output instance of RichTextLines.
236
237    Raises:
238      ValueError: If begin or end is negative.
239    """
240
241    if begin < 0 or end < 0:
242      raise ValueError("Encountered negative index.")
243
244    # Copy lines.
245    lines = self.lines[begin:end]
246
247    # Slice font attribute segments.
248    font_attr_segs = {}
249    for key in self.font_attr_segs:
250      if key >= begin and key < end:
251        font_attr_segs[key - begin] = self.font_attr_segs[key]
252
253    # Slice annotations.
254    annotations = {}
255    for key in self.annotations:
256      if not isinstance(key, int):
257        # Annotations can contain keys that are not line numbers.
258        annotations[key] = self.annotations[key]
259      elif key >= begin and key < end:
260        annotations[key - begin] = self.annotations[key]
261
262    return RichTextLines(
263        lines, font_attr_segs=font_attr_segs, annotations=annotations)
264
265  def extend(self, other):
266    """Extend this instance of RichTextLines with another instance.
267
268    The extension takes effect on the text lines, the font attribute segments,
269    as well as the annotations. The line indices in the font attribute
270    segments and the annotations are adjusted to account for the existing
271    lines. If there are duplicate, non-line-index fields in the annotations,
272    the value from the input argument "other" will override that in this
273    instance.
274
275    Args:
276      other: (RichTextLines) The other RichTextLines instance to be appended at
277        the end of this instance.
278    """
279
280    orig_num_lines = self.num_lines()  # Record original number of lines.
281
282    # Merge the lines.
283    self._lines.extend(other.lines)
284
285    # Merge the font_attr_segs.
286    for line_index in other.font_attr_segs:
287      self._font_attr_segs[orig_num_lines + line_index] = (
288          other.font_attr_segs[line_index])
289
290    # Merge the annotations.
291    for key in other.annotations:
292      if isinstance(key, int):
293        self._annotations[orig_num_lines + key] = (other.annotations[key])
294      else:
295        self._annotations[key] = other.annotations[key]
296
297  def _extend_before(self, other):
298    """Add another RichTextLines object to the front.
299
300    Args:
301      other: (RichTextLines) The other object to add to the front to this
302        object.
303    """
304
305    other_num_lines = other.num_lines()  # Record original number of lines.
306
307    # Merge the lines.
308    self._lines = other.lines + self._lines
309
310    # Merge the font_attr_segs.
311    new_font_attr_segs = {}
312    for line_index in self.font_attr_segs:
313      new_font_attr_segs[other_num_lines + line_index] = (
314          self.font_attr_segs[line_index])
315    new_font_attr_segs.update(other.font_attr_segs)
316    self._font_attr_segs = new_font_attr_segs
317
318    # Merge the annotations.
319    new_annotations = {}
320    for key in self._annotations:
321      if isinstance(key, int):
322        new_annotations[other_num_lines + key] = (self.annotations[key])
323      else:
324        new_annotations[key] = other.annotations[key]
325
326    new_annotations.update(other.annotations)
327    self._annotations = new_annotations
328
329  def append(self, line, font_attr_segs=None):
330    """Append a single line of text.
331
332    Args:
333      line: (str) The text to be added to the end.
334      font_attr_segs: (list of tuples) Font attribute segments of the appended
335        line.
336    """
337
338    self._lines.append(line)
339    if font_attr_segs:
340      self._font_attr_segs[len(self._lines) - 1] = font_attr_segs
341
342  def append_rich_line(self, rich_line):
343    self.append(rich_line.text, rich_line.font_attr_segs)
344
345  def prepend(self, line, font_attr_segs=None):
346    """Prepend (i.e., add to the front) a single line of text.
347
348    Args:
349      line: (str) The text to be added to the front.
350      font_attr_segs: (list of tuples) Font attribute segments of the appended
351        line.
352    """
353
354    other = RichTextLines(line)
355    if font_attr_segs:
356      other.font_attr_segs[0] = font_attr_segs
357    self._extend_before(other)
358
359  def write_to_file(self, file_path):
360    """Write the object itself to file, in a plain format.
361
362    The font_attr_segs and annotations are ignored.
363
364    Args:
365      file_path: (str) path of the file to write to.
366    """
367
368    with gfile.Open(file_path, "w") as f:
369      for line in self._lines:
370        f.write(line + "\n")
371
372  # TODO(cais): Add a method to allow appending to a line in RichTextLines with
373  # both text and font_attr_segs.
374
375
376def regex_find(orig_screen_output, regex, font_attr):
377  """Perform regex match in rich text lines.
378
379  Produces a new RichTextLines object with font_attr_segs containing highlighted
380  regex matches.
381
382  Example use cases include:
383  1) search for specific items in a large list of items, and
384  2) search for specific numerical values in a large tensor.
385
386  Args:
387    orig_screen_output: The original RichTextLines, in which the regex find
388      is to be performed.
389    regex: The regex used for matching.
390    font_attr: Font attribute used for highlighting the found result.
391
392  Returns:
393    A modified copy of orig_screen_output.
394
395  Raises:
396    ValueError: If input str regex is not a valid regular expression.
397  """
398  new_screen_output = RichTextLines(
399      orig_screen_output.lines,
400      font_attr_segs=copy.deepcopy(orig_screen_output.font_attr_segs),
401      annotations=orig_screen_output.annotations)
402
403  try:
404    re_prog = re.compile(regex)
405  except sre_constants.error:
406    raise ValueError("Invalid regular expression: \"%s\"" % regex)
407
408  regex_match_lines = []
409  for i, line in enumerate(new_screen_output.lines):
410    find_it = re_prog.finditer(line)
411
412    match_segs = []
413    for match in find_it:
414      match_segs.append((match.start(), match.end(), font_attr))
415
416    if match_segs:
417      if i not in new_screen_output.font_attr_segs:
418        new_screen_output.font_attr_segs[i] = match_segs
419      else:
420        new_screen_output.font_attr_segs[i].extend(match_segs)
421        new_screen_output.font_attr_segs[i] = sorted(
422            new_screen_output.font_attr_segs[i], key=lambda x: x[0])
423      regex_match_lines.append(i)
424
425  new_screen_output.annotations[REGEX_MATCH_LINES_KEY] = regex_match_lines
426  return new_screen_output
427
428
429def wrap_rich_text_lines(inp, cols):
430  """Wrap RichTextLines according to maximum number of columns.
431
432  Produces a new RichTextLines object with the text lines, font_attr_segs and
433  annotations properly wrapped. This ought to be used sparingly, as in most
434  cases, command handlers producing RichTextLines outputs should know the
435  screen/panel width via the screen_info kwarg and should produce properly
436  length-limited lines in the output accordingly.
437
438  Args:
439    inp: Input RichTextLines object.
440    cols: Number of columns, as an int.
441
442  Returns:
443    1) A new instance of RichTextLines, with line lengths limited to cols.
444    2) A list of new (wrapped) line index. For example, if the original input
445      consists of three lines and only the second line is wrapped, and it's
446      wrapped into two lines, this return value will be: [0, 1, 3].
447  Raises:
448    ValueError: If inputs have invalid types.
449  """
450
451  new_line_indices = []
452
453  if not isinstance(inp, RichTextLines):
454    raise ValueError("Invalid type of input screen_output")
455
456  if not isinstance(cols, int):
457    raise ValueError("Invalid type of input cols")
458
459  out = RichTextLines([])
460
461  row_counter = 0  # Counter for new row index
462  for i, line in enumerate(inp.lines):
463    new_line_indices.append(out.num_lines())
464
465    if i in inp.annotations:
466      out.annotations[row_counter] = inp.annotations[i]
467
468    if len(line) <= cols:
469      # No wrapping.
470      out.lines.append(line)
471      if i in inp.font_attr_segs:
472        out.font_attr_segs[row_counter] = inp.font_attr_segs[i]
473
474      row_counter += 1
475    else:
476      # Wrap.
477      wlines = []  # Wrapped lines.
478
479      osegs = []
480      if i in inp.font_attr_segs:
481        osegs = inp.font_attr_segs[i]
482
483      idx = 0
484      while idx < len(line):
485        if idx + cols > len(line):
486          rlim = len(line)
487        else:
488          rlim = idx + cols
489
490        wlines.append(line[idx:rlim])
491        for seg in osegs:
492          if (seg[0] < rlim) and (seg[1] >= idx):
493            # Calculate left bound within wrapped line.
494            if seg[0] >= idx:
495              lb = seg[0] - idx
496            else:
497              lb = 0
498
499            # Calculate right bound within wrapped line.
500            if seg[1] < rlim:
501              rb = seg[1] - idx
502            else:
503              rb = rlim - idx
504
505            if rb > lb:  # Omit zero-length segments.
506              wseg = (lb, rb, seg[2])
507              if row_counter not in out.font_attr_segs:
508                out.font_attr_segs[row_counter] = [wseg]
509              else:
510                out.font_attr_segs[row_counter].append(wseg)
511
512        idx += cols
513        row_counter += 1
514
515      out.lines.extend(wlines)
516
517  # Copy over keys of annotation that are not row indices.
518  for key in inp.annotations:
519    if not isinstance(key, int):
520      out.annotations[key] = inp.annotations[key]
521
522  return out, new_line_indices
523
524
525class CommandHandlerRegistry:
526  """Registry of command handlers for CLI.
527
528  Handler methods (callables) for user commands can be registered with this
529  class, which then is able to dispatch commands to the correct handlers and
530  retrieve the RichTextLines output.
531
532  For example, suppose you have the following handler defined:
533    def echo(argv, screen_info=None):
534      return RichTextLines(["arguments = %s" % " ".join(argv),
535                            "screen_info = " + repr(screen_info)])
536
537  you can register the handler with the command prefix "echo" and alias "e":
538    registry = CommandHandlerRegistry()
539    registry.register_command_handler("echo", echo,
540        "Echo arguments, along with screen info", prefix_aliases=["e"])
541
542  then to invoke this command handler with some arguments and screen_info, do:
543    registry.dispatch_command("echo", ["foo", "bar"], screen_info={"cols": 80})
544
545  or with the prefix alias:
546    registry.dispatch_command("e", ["foo", "bar"], screen_info={"cols": 80})
547
548  The call will return a RichTextLines object which can be rendered by a CLI.
549  """
550
551  HELP_COMMAND = "help"
552  HELP_COMMAND_ALIASES = ["h"]
553  VERSION_COMMAND = "version"
554  VERSION_COMMAND_ALIASES = ["ver"]
555
556  def __init__(self):
557    # A dictionary from command prefix to handler.
558    self._handlers = {}
559
560    # A dictionary from prefix alias to prefix.
561    self._alias_to_prefix = {}
562
563    # A dictionary from prefix to aliases.
564    self._prefix_to_aliases = {}
565
566    # A dictionary from command prefix to help string.
567    self._prefix_to_help = {}
568
569    # Introductory text to help information.
570    self._help_intro = None
571
572    # Register a default handler for the command "help".
573    self.register_command_handler(
574        self.HELP_COMMAND,
575        self._help_handler,
576        "Print this help message.",
577        prefix_aliases=self.HELP_COMMAND_ALIASES)
578
579    # Register a default handler for the command "version".
580    self.register_command_handler(
581        self.VERSION_COMMAND,
582        self._version_handler,
583        "Print the versions of TensorFlow and its key dependencies.",
584        prefix_aliases=self.VERSION_COMMAND_ALIASES)
585
586  def register_command_handler(self,
587                               prefix,
588                               handler,
589                               help_info,
590                               prefix_aliases=None):
591    """Register a callable as a command handler.
592
593    Args:
594      prefix: Command prefix, i.e., the first word in a command, e.g.,
595        "print" as in "print tensor_1".
596      handler: A callable of the following signature:
597          foo_handler(argv, screen_info=None),
598        where argv is the argument vector (excluding the command prefix) and
599          screen_info is a dictionary containing information about the screen,
600          such as number of columns, e.g., {"cols": 100}.
601        The callable should return:
602          1) a RichTextLines object representing the screen output.
603
604        The callable can also raise an exception of the type CommandLineExit,
605        which if caught by the command-line interface, will lead to its exit.
606        The exception can optionally carry an exit token of arbitrary type.
607      help_info: A help string.
608      prefix_aliases: Aliases for the command prefix, as a list of str. E.g.,
609        shorthands for the command prefix: ["p", "pr"]
610
611    Raises:
612      ValueError: If
613        1) the prefix is empty, or
614        2) handler is not callable, or
615        3) a handler is already registered for the prefix, or
616        4) elements in prefix_aliases clash with existing aliases.
617        5) help_info is not a str.
618    """
619
620    if not prefix:
621      raise ValueError("Empty command prefix")
622
623    if prefix in self._handlers:
624      raise ValueError(
625          "A handler is already registered for command prefix \"%s\"" % prefix)
626
627    # Make sure handler is callable.
628    if not callable(handler):
629      raise ValueError("handler is not callable")
630
631    # Make sure that help info is a string.
632    if not isinstance(help_info, str):
633      raise ValueError("help_info is not a str")
634
635    # Process prefix aliases.
636    if prefix_aliases:
637      for alias in prefix_aliases:
638        if self._resolve_prefix(alias):
639          raise ValueError(
640              "The prefix alias \"%s\" clashes with existing prefixes or "
641              "aliases." % alias)
642        self._alias_to_prefix[alias] = prefix
643
644      self._prefix_to_aliases[prefix] = prefix_aliases
645
646    # Store handler.
647    self._handlers[prefix] = handler
648
649    # Store help info.
650    self._prefix_to_help[prefix] = help_info
651
652  def dispatch_command(self, prefix, argv, screen_info=None):
653    """Handles a command by dispatching it to a registered command handler.
654
655    Args:
656      prefix: Command prefix, as a str, e.g., "print".
657      argv: Command argument vector, excluding the command prefix, represented
658        as a list of str, e.g.,
659        ["tensor_1"]
660      screen_info: A dictionary containing screen info, e.g., {"cols": 100}.
661
662    Returns:
663      An instance of RichTextLines or None. If any exception is caught during
664      the invocation of the command handler, the RichTextLines will wrap the
665      error type and message.
666
667    Raises:
668      ValueError: If
669        1) prefix is empty, or
670        2) no command handler is registered for the command prefix, or
671        3) the handler is found for the prefix, but it fails to return a
672          RichTextLines or raise any exception.
673      CommandLineExit:
674        If the command handler raises this type of exception, this method will
675        simply pass it along.
676    """
677    if not prefix:
678      raise ValueError("Prefix is empty")
679
680    resolved_prefix = self._resolve_prefix(prefix)
681    if not resolved_prefix:
682      raise ValueError("No handler is registered for command prefix \"%s\"" %
683                       prefix)
684
685    handler = self._handlers[resolved_prefix]
686    try:
687      output = handler(argv, screen_info=screen_info)
688    except CommandLineExit as e:
689      raise e
690    except SystemExit as e:
691      # Special case for syntax errors caught by argparse.
692      lines = ["Syntax error for command: %s" % prefix,
693               "For help, do \"help %s\"" % prefix]
694      output = RichTextLines(lines)
695
696    except BaseException as e:  # pylint: disable=broad-except
697      lines = ["Error occurred during handling of command: %s %s:" %
698               (resolved_prefix, " ".join(argv)), "%s: %s" % (type(e), str(e))]
699
700      # Include traceback of the exception.
701      lines.append("")
702      lines.extend(traceback.format_exc().split("\n"))
703
704      output = RichTextLines(lines)
705
706    if not isinstance(output, RichTextLines) and output is not None:
707      raise ValueError(
708          "Return value from command handler %s is not None or a RichTextLines "
709          "instance" % str(handler))
710
711    return output
712
713  def is_registered(self, prefix):
714    """Test if a command prefix or its alias is has a registered handler.
715
716    Args:
717      prefix: A prefix or its alias, as a str.
718
719    Returns:
720      True iff a handler is registered for prefix.
721    """
722    return self._resolve_prefix(prefix) is not None
723
724  def get_help(self, cmd_prefix=None):
725    """Compile help information into a RichTextLines object.
726
727    Args:
728      cmd_prefix: Optional command prefix. As the prefix itself or one of its
729        aliases.
730
731    Returns:
732      A RichTextLines object containing the help information. If cmd_prefix
733      is None, the return value will be the full command-line help. Otherwise,
734      it will be the help information for the specified command.
735    """
736    if not cmd_prefix:
737      # Print full help information, in sorted order of the command prefixes.
738      help_info = RichTextLines([])
739      if self._help_intro:
740        # If help intro is available, show it at the beginning.
741        help_info.extend(self._help_intro)
742
743      sorted_prefixes = sorted(self._handlers)
744      for cmd_prefix in sorted_prefixes:
745        lines = self._get_help_for_command_prefix(cmd_prefix)
746        lines.append("")
747        lines.append("")
748        help_info.extend(RichTextLines(lines))
749
750      return help_info
751    else:
752      return RichTextLines(self._get_help_for_command_prefix(cmd_prefix))
753
754  def set_help_intro(self, help_intro):
755    """Set an introductory message to help output.
756
757    Args:
758      help_intro: (RichTextLines) Rich text lines appended to the
759        beginning of the output of the command "help", as introductory
760        information.
761    """
762    self._help_intro = help_intro
763
764  def _help_handler(self, args, screen_info=None):
765    """Command handler for "help".
766
767    "help" is a common command that merits built-in support from this class.
768
769    Args:
770      args: Command line arguments to "help" (not including "help" itself).
771      screen_info: (dict) Information regarding the screen, e.g., the screen
772        width in characters: {"cols": 80}
773
774    Returns:
775      (RichTextLines) Screen text output.
776    """
777
778    _ = screen_info  # Unused currently.
779
780    if not args:
781      return self.get_help()
782    elif len(args) == 1:
783      return self.get_help(args[0])
784    else:
785      return RichTextLines(["ERROR: help takes only 0 or 1 input argument."])
786
787  def _version_handler(self, args, screen_info=None):
788    del args  # Unused currently.
789    del screen_info  # Unused currently.
790    return get_tensorflow_version_lines(include_dependency_versions=True)
791
792  def _resolve_prefix(self, token):
793    """Resolve command prefix from the prefix itself or its alias.
794
795    Args:
796      token: a str to be resolved.
797
798    Returns:
799      If resolvable, the resolved command prefix.
800      If not resolvable, None.
801    """
802    if token in self._handlers:
803      return token
804    elif token in self._alias_to_prefix:
805      return self._alias_to_prefix[token]
806    else:
807      return None
808
809  def _get_help_for_command_prefix(self, cmd_prefix):
810    """Compile the help information for a given command prefix.
811
812    Args:
813      cmd_prefix: Command prefix, as the prefix itself or one of its
814        aliases.
815
816    Returns:
817      A list of str as the help information fo cmd_prefix. If the cmd_prefix
818        does not exist, the returned list of str will indicate that.
819    """
820    lines = []
821
822    resolved_prefix = self._resolve_prefix(cmd_prefix)
823    if not resolved_prefix:
824      lines.append("Invalid command prefix: \"%s\"" % cmd_prefix)
825      return lines
826
827    lines.append(resolved_prefix)
828
829    if resolved_prefix in self._prefix_to_aliases:
830      lines.append(HELP_INDENT + "Aliases: " + ", ".join(
831          self._prefix_to_aliases[resolved_prefix]))
832
833    lines.append("")
834    help_lines = self._prefix_to_help[resolved_prefix].split("\n")
835    for line in help_lines:
836      lines.append(HELP_INDENT + line)
837
838    return lines
839
840
841class TabCompletionRegistry:
842  """Registry for tab completion responses."""
843
844  def __init__(self):
845    self._comp_dict = {}
846
847  # TODO(cais): Rename method names with "comp" to "*completion*" to avoid
848  # confusion.
849
850  def register_tab_comp_context(self, context_words, comp_items):
851    """Register a tab-completion context.
852
853    Register that, for each word in context_words, the potential tab-completions
854    are the words in comp_items.
855
856    A context word is a pre-existing, completed word in the command line that
857    determines how tab-completion works for another, incomplete word in the same
858    command line.
859    Completion items consist of potential candidates for the incomplete word.
860
861    To give a general example, a context word can be "drink", and the completion
862    items can be ["coffee", "tea", "water"]
863
864    Note: A context word can be empty, in which case the context is for the
865     top-level commands.
866
867    Args:
868      context_words: A list of context words belonging to the context being
869        registered. It is a list of str, instead of a single string, to support
870        synonym words triggering the same tab-completion context, e.g.,
871        both "drink" and the short-hand "dr" can trigger the same context.
872      comp_items: A list of completion items, as a list of str.
873
874    Raises:
875      TypeError: if the input arguments are not all of the correct types.
876    """
877
878    if not isinstance(context_words, list):
879      raise TypeError("Incorrect type in context_list: Expected list, got %s" %
880                      type(context_words))
881
882    if not isinstance(comp_items, list):
883      raise TypeError("Incorrect type in comp_items: Expected list, got %s" %
884                      type(comp_items))
885
886    # Sort the completion items on registration, so that later during
887    # get_completions calls, no sorting will be necessary.
888    sorted_comp_items = sorted(comp_items)
889
890    for context_word in context_words:
891      self._comp_dict[context_word] = sorted_comp_items
892
893  def deregister_context(self, context_words):
894    """Deregister a list of context words.
895
896    Args:
897      context_words: A list of context words to deregister, as a list of str.
898
899    Raises:
900      KeyError: if there are word(s) in context_words that do not correspond
901        to any registered contexts.
902    """
903
904    for context_word in context_words:
905      if context_word not in self._comp_dict:
906        raise KeyError("Cannot deregister unregistered context word \"%s\"" %
907                       context_word)
908
909    for context_word in context_words:
910      del self._comp_dict[context_word]
911
912  def extend_comp_items(self, context_word, new_comp_items):
913    """Add a list of completion items to a completion context.
914
915    Args:
916      context_word: A single completion word as a string. The extension will
917        also apply to all other context words of the same context.
918      new_comp_items: (list of str) New completion items to add.
919
920    Raises:
921      KeyError: if the context word has not been registered.
922    """
923
924    if context_word not in self._comp_dict:
925      raise KeyError("Context word \"%s\" has not been registered" %
926                     context_word)
927
928    self._comp_dict[context_word].extend(new_comp_items)
929    self._comp_dict[context_word] = sorted(self._comp_dict[context_word])
930
931  def remove_comp_items(self, context_word, comp_items):
932    """Remove a list of completion items from a completion context.
933
934    Args:
935      context_word: A single completion word as a string. The removal will
936        also apply to all other context words of the same context.
937      comp_items: Completion items to remove.
938
939    Raises:
940      KeyError: if the context word has not been registered.
941    """
942
943    if context_word not in self._comp_dict:
944      raise KeyError("Context word \"%s\" has not been registered" %
945                     context_word)
946
947    for item in comp_items:
948      self._comp_dict[context_word].remove(item)
949
950  def get_completions(self, context_word, prefix):
951    """Get the tab completions given a context word and a prefix.
952
953    Args:
954      context_word: The context word.
955      prefix: The prefix of the incomplete word.
956
957    Returns:
958      (1) None if no registered context matches the context_word.
959          A list of str for the matching completion items. Can be an empty list
960          of a matching context exists, but no completion item matches the
961          prefix.
962      (2) Common prefix of all the words in the first return value. If the
963          first return value is None, this return value will be None, too. If
964          the first return value is not None, i.e., a list, this return value
965          will be a str, which can be an empty str if there is no common
966          prefix among the items of the list.
967    """
968
969    if context_word not in self._comp_dict:
970      return None, None
971
972    comp_items = self._comp_dict[context_word]
973    comp_items = sorted(
974        [item for item in comp_items if item.startswith(prefix)])
975
976    return comp_items, self._common_prefix(comp_items)
977
978  def _common_prefix(self, m):
979    """Given a list of str, returns the longest common prefix.
980
981    Args:
982      m: (list of str) A list of strings.
983
984    Returns:
985      (str) The longest common prefix.
986    """
987    if not m:
988      return ""
989
990    s1 = min(m)
991    s2 = max(m)
992    for i, c in enumerate(s1):
993      if c != s2[i]:
994        return s1[:i]
995
996    return s1
997
998
999class CommandHistory:
1000  """Keeps command history and supports lookup."""
1001
1002  _HISTORY_FILE_NAME = ".tfdbg_history"
1003
1004  def __init__(self, limit=100, history_file_path=None):
1005    """CommandHistory constructor.
1006
1007    Args:
1008      limit: Maximum number of the most recent commands that this instance
1009        keeps track of, as an int.
1010      history_file_path: (str) Manually specified path to history file. Used in
1011        testing.
1012    """
1013
1014    self._commands = []
1015    self._limit = limit
1016    self._history_file_path = (
1017        history_file_path or self._get_default_history_file_path())
1018    self._load_history_from_file()
1019
1020  def _load_history_from_file(self):
1021    if os.path.isfile(self._history_file_path):
1022      try:
1023        with open(self._history_file_path, "rt") as history_file:
1024          commands = history_file.readlines()
1025        self._commands = [command.strip() for command in commands
1026                          if command.strip()]
1027
1028        # Limit the size of the history file.
1029        if len(self._commands) > self._limit:
1030          self._commands = self._commands[-self._limit:]
1031          with open(self._history_file_path, "wt") as history_file:
1032            for command in self._commands:
1033              history_file.write(command + "\n")
1034      except IOError:
1035        print("WARNING: writing history file failed.")
1036
1037  def _add_command_to_history_file(self, command):
1038    try:
1039      with open(self._history_file_path, "at") as history_file:
1040        history_file.write(command + "\n")
1041    except IOError:
1042      pass
1043
1044  @classmethod
1045  def _get_default_history_file_path(cls):
1046    return os.path.join(os.path.expanduser("~"), cls._HISTORY_FILE_NAME)
1047
1048  def add_command(self, command):
1049    """Add a command to the command history.
1050
1051    Args:
1052      command: The history command, as a str.
1053
1054    Raises:
1055      TypeError: if command is not a str.
1056    """
1057
1058    if self._commands and command == self._commands[-1]:
1059      # Ignore repeating commands in a row.
1060      return
1061
1062    if not isinstance(command, str):
1063      raise TypeError("Attempt to enter non-str entry to command history")
1064
1065    self._commands.append(command)
1066
1067    if len(self._commands) > self._limit:
1068      self._commands = self._commands[-self._limit:]
1069
1070    self._add_command_to_history_file(command)
1071
1072  def most_recent_n(self, n):
1073    """Look up the n most recent commands.
1074
1075    Args:
1076      n: Number of most recent commands to look up.
1077
1078    Returns:
1079      A list of n most recent commands, or all available most recent commands,
1080      if n exceeds size of the command history, in chronological order.
1081    """
1082
1083    return self._commands[-n:]
1084
1085  def lookup_prefix(self, prefix, n):
1086    """Look up the n most recent commands that starts with prefix.
1087
1088    Args:
1089      prefix: The prefix to lookup.
1090      n: Number of most recent commands to look up.
1091
1092    Returns:
1093      A list of n most recent commands that have the specified prefix, or all
1094      available most recent commands that have the prefix, if n exceeds the
1095      number of history commands with the prefix.
1096    """
1097
1098    commands = [cmd for cmd in self._commands if cmd.startswith(prefix)]
1099
1100    return commands[-n:]
1101
1102  # TODO(cais): Lookup by regex.
1103
1104
1105class MenuItem:
1106  """A class for an item in a text-based menu."""
1107
1108  def __init__(self, caption, content, enabled=True):
1109    """Menu constructor.
1110
1111    TODO(cais): Nested menu is currently not supported. Support it.
1112
1113    Args:
1114      caption: (str) caption of the menu item.
1115      content: Content of the menu item. For a menu item that triggers
1116        a command, for example, content is the command string.
1117      enabled: (bool) whether this menu item is enabled.
1118    """
1119
1120    self._caption = caption
1121    self._content = content
1122    self._enabled = enabled
1123
1124  @property
1125  def caption(self):
1126    return self._caption
1127
1128  @property
1129  def type(self):
1130    return self._node_type
1131
1132  @property
1133  def content(self):
1134    return self._content
1135
1136  def is_enabled(self):
1137    return self._enabled
1138
1139  def disable(self):
1140    self._enabled = False
1141
1142  def enable(self):
1143    self._enabled = True
1144
1145
1146class Menu:
1147  """A class for text-based menu."""
1148
1149  def __init__(self, name=None):
1150    """Menu constructor.
1151
1152    Args:
1153      name: (str or None) name of this menu.
1154    """
1155
1156    self._name = name
1157    self._items = []
1158
1159  def append(self, item):
1160    """Append an item to the Menu.
1161
1162    Args:
1163      item: (MenuItem) the item to be appended.
1164    """
1165    self._items.append(item)
1166
1167  def insert(self, index, item):
1168    self._items.insert(index, item)
1169
1170  def num_items(self):
1171    return len(self._items)
1172
1173  def captions(self):
1174    return [item.caption for item in self._items]
1175
1176  def caption_to_item(self, caption):
1177    """Get a MenuItem from the caption.
1178
1179    Args:
1180      caption: (str) The caption to look up.
1181
1182    Returns:
1183      (MenuItem) The first-match menu item with the caption, if any.
1184
1185    Raises:
1186      LookupError: If a menu item with the caption does not exist.
1187    """
1188
1189    captions = self.captions()
1190    if caption not in captions:
1191      raise LookupError("There is no menu item with the caption \"%s\"" %
1192                        caption)
1193
1194    return self._items[captions.index(caption)]
1195
1196  def format_as_single_line(self,
1197                            prefix=None,
1198                            divider=" | ",
1199                            enabled_item_attrs=None,
1200                            disabled_item_attrs=None):
1201    """Format the menu as a single-line RichTextLines object.
1202
1203    Args:
1204      prefix: (str) String added to the beginning of the line.
1205      divider: (str) The dividing string between the menu items.
1206      enabled_item_attrs: (list or str) Attributes applied to each enabled
1207        menu item, e.g., ["bold", "underline"].
1208      disabled_item_attrs: (list or str) Attributes applied to each
1209        disabled menu item, e.g., ["red"].
1210
1211    Returns:
1212      (RichTextLines) A single-line output representing the menu, with
1213        font_attr_segs marking the individual menu items.
1214    """
1215
1216    if (enabled_item_attrs is not None and
1217        not isinstance(enabled_item_attrs, list)):
1218      enabled_item_attrs = [enabled_item_attrs]
1219
1220    if (disabled_item_attrs is not None and
1221        not isinstance(disabled_item_attrs, list)):
1222      disabled_item_attrs = [disabled_item_attrs]
1223
1224    menu_line = prefix if prefix is not None else ""
1225    attr_segs = []
1226
1227    for item in self._items:
1228      menu_line += item.caption
1229      item_name_begin = len(menu_line) - len(item.caption)
1230
1231      if item.is_enabled():
1232        final_attrs = [item]
1233        if enabled_item_attrs:
1234          final_attrs.extend(enabled_item_attrs)
1235        attr_segs.append((item_name_begin, len(menu_line), final_attrs))
1236      else:
1237        if disabled_item_attrs:
1238          attr_segs.append(
1239              (item_name_begin, len(menu_line), disabled_item_attrs))
1240
1241      menu_line += divider
1242
1243    return RichTextLines(menu_line, font_attr_segs={0: attr_segs})
1244