Input:
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Building Blocks of TensorFlow Debugger Command-Line Interface."""
import copy
import os
import re
import traceback

import numpy as np

from tensorflow.python.client import pywrap_tf_session
from tensorflow.python.platform import gfile

HELP_INDENT = "  "

EXPLICIT_USER_EXIT = "explicit_user_exit"
REGEX_MATCH_LINES_KEY = "regex_match_lines"
INIT_SCROLL_POS_KEY = "init_scroll_pos"

MAIN_MENU_KEY = "mm:"

class CommandLineExit(Exception):

  def __init__(self, exit_token=None):
    Exception.__init__(self)
    self._exit_token = exit_token

  @property
  def exit_token(self):
    return self._exit_token

class RichLine:
  """Rich single-line text.

  Attributes:
    text: A plain string, the raw text represented by this object.  Should not
      contain newlines.
    font_attr_segs: A list of (start, end, font attribute) triples, representing
      richness information applied to substrings of text.
  """

  def __init__(self, text="", font_attr=None):
    """Construct a RichLine with no rich attributes or a single attribute.

    Args:
      text: Raw text string
      font_attr: If specified, a single font attribute to be applied to the
        entire text.  Extending this object via concatenation allows creation
        of text with varying attributes.
    """
    # TODO(ebreck) Make .text and .font_attr protected members when we no
    # longer need public access.
    self.text = text
    if font_attr:
      self.font_attr_segs = [(0, len(text), font_attr)]
    else:
      self.font_attr_segs = []

  def __add__(self, other):
    """Concatenate two chunks of maybe rich text to make a longer rich line.

    Does not modify self.

    Args:
      other: Another piece of text to concatenate with this one.
        If it is a plain str, it will be appended to this string with no
        attributes.  If it is a RichLine, it will be appended to this string
        with its attributes preserved.

    Returns:
      A new RichLine comprising both chunks of text, with appropriate
        attributes applied to the corresponding substrings.
    """
    ret = RichLine()
    if isinstance(other, str):
      ret.text = self.text + other
      ret.font_attr_segs = self.font_attr_segs[:]
      return ret
    elif isinstance(other, RichLine):
      ret.text = self.text + other.text
      ret.font_attr_segs = self.font_attr_segs[:]
      old_len = len(self.text)
      for start, end, font_attr in other.font_attr_segs:
        ret.font_attr_segs.append((old_len + start, old_len + end, font_attr))
      return ret
    else:
      raise TypeError("%r cannot be concatenated with a RichLine" % other)

  def __len__(self):
    return len(self.text)

def rich_text_lines_from_rich_line_list(rich_text_list, annotations=None):
  """Convert a list of RichLine objects or strings to a RichTextLines object.

  Args:
    rich_text_list: a list of RichLine objects or strings
    annotations: annotations for the resultant RichTextLines object.

  Returns:
    A corresponding RichTextLines object.
  """
  lines = []
  font_attr_segs = {}
  for i, rl in enumerate(rich_text_list):
    if isinstance(rl, RichLine):
      lines.append(rl.text)
      if rl.font_attr_segs:
        font_attr_segs[i] = rl.font_attr_segs
    else:
      lines.append(rl)
  return RichTextLines(lines, font_attr_segs, annotations=annotations)

def get_tensorflow_version_lines(include_dependency_versions=False):
  """Generate RichTextLines with TensorFlow version info.

  Args:
    include_dependency_versions: Include the version of TensorFlow's key
      dependencies, such as numpy.

  Returns:
    A formatted, multi-line `RichTextLines` object.
  """
  lines = ["TensorFlow version: %s" % pywrap_tf_session.__version__]
  lines.append("")
  if include_dependency_versions:
    lines.append("Dependency version(s):")
    lines.append("  numpy: %s" % np.__version__)
    lines.append("")
  return RichTextLines(lines)

class RichTextLines:
  """Rich multi-line text.

  Line-by-line text output, with font attributes (e.g., color) and annotations
  (e.g., indices in a multi-dimensional tensor). Used as the text output of CLI
  commands. Can be rendered on terminal environments such as curses.

  This is not to be confused with Rich Text Format (RTF). This class is for text
  lines only.
  """

  def __init__(self, lines, font_attr_segs=None, annotations=None):
    """Constructor of RichTextLines.

    Args:
      lines: A list of str or a single str, representing text output to
        screen. The latter case is for convenience when the text output is
        single-line.
      font_attr_segs: A map from 0-based row index to a list of 3-tuples.
        It lists segments in each row that have special font attributes, such
        as colors, that are not the default attribute. For example:
        {1: [(0, 3, "red"), (4, 7, "green")], 2: [(10, 20, "yellow")]}

        In each tuple, the 1st element is the start index of the segment. The
        2nd element is the end index, in an "open interval" fashion. The 3rd
        element is an object or a list of objects that represents the font
        attribute. Colors are represented as strings as in the examples above.
      annotations: A map from 0-based row index to any object for annotating
        the row. A typical use example is annotating rows of the output as
        indices in a multi-dimensional tensor. For example, consider the
        following text representation of a 3x2x2 tensor:
          [[[0, 0], [0, 0]],
           [[0, 0], [0, 0]],
           [[0, 0], [0, 0]]]
        The annotation can indicate the indices of the first element shown in
        each row, i.e.,
          {0: [0, 0, 0], 1: [1, 0, 0], 2: [2, 0, 0]}
        This information can make display of tensors on screen clearer and can
        help the user navigate (scroll) to the desired location in a large
        tensor.

    Raises:
      ValueError: If lines is of invalid type.
    """
    if isinstance(lines, list):
      self._lines = lines
    elif isinstance(lines, str):
      self._lines = [lines]
    else:
      raise ValueError("Unexpected type in lines: %s" % type(lines))

    self._font_attr_segs = font_attr_segs
    if not self._font_attr_segs:
      self._font_attr_segs = {}
      # TODO(cais): Refactor to collections.defaultdict(list) to simplify code.

    self._annotations = annotations
    if not self._annotations:
      self._annotations = {}
      # TODO(cais): Refactor to collections.defaultdict(list) to simplify code.

  @property
  def lines(self):
    return self._lines

  @property
  def font_attr_segs(self):
    return self._font_attr_segs

  @property
  def annotations(self):
    return self._annotations

  def num_lines(self):
    return len(self._lines)

  def slice(self, begin, end):
    """Slice a RichTextLines object.

    The object itself is not changed. A sliced instance is returned.

    Args:
      begin: (int) Beginning line index (inclusive). Must be >= 0.
      end: (int) Ending line index (exclusive). Must be >= 0.

    Returns:
      (RichTextLines) Sliced output instance of RichTextLines.

    Raises:
      ValueError: If begin or end is negative.
    """

    if begin < 0 or end < 0:
      raise ValueError("Encountered negative index.")

    # Copy lines.
    lines = self.lines[begin:end]

    # Slice font attribute segments.
    font_attr_segs = {}
    for key in self.font_attr_segs:
      if key >= begin and key < end:
        font_attr_segs[key - begin] = self.font_attr_segs[key]

    # Slice annotations.
    annotations = {}
    for key in self.annotations:
      if not isinstance(key, int):
        # Annotations can contain keys that are not line numbers.
        annotations[key] = self.annotations[key]
      elif key >= begin and key < end:
        annotations[key - begin] = self.annotations[key]

    return RichTextLines(
        lines, font_attr_segs=font_attr_segs, annotations=annotations)

  def extend(self, other):
    """Extend this instance of RichTextLines with another instance.

    The extension takes effect on the text lines, the font attribute segments,
    as well as the annotations. The line indices in the font attribute
    segments and the annotations are adjusted to account for the existing
    lines. If there are duplicate, non-line-index fields in the annotations,
    the value from the input argument "other" will override that in this
    instance.

    Args:
      other: (RichTextLines) The other RichTextLines instance to be appended at
        the end of this instance.
    """

    orig_num_lines = self.num_lines()  # Record original number of lines.

    # Merge the lines.
    self._lines.extend(other.lines)

    # Merge the font_attr_segs.
    for line_index in other.font_attr_segs:
      self._font_attr_segs[orig_num_lines + line_index] = (
          other.font_attr_segs[line_index])

    # Merge the annotations.
    for key in other.annotations:
      if isinstance(key, int):
        self._annotations[orig_num_lines + key] = (other.annotations[key])
      else:
        self._annotations[key] = other.annotations[key]

  def _extend_before(self, other):
    """Add another RichTextLines object to the front.

    Args:
      other: (RichTextLines) The other object to add to the front to this
        object.
    """

    other_num_lines = other.num_lines()  # Record original number of lines.

    # Merge the lines.
    self._lines = other.lines + self._lines

    # Merge the font_attr_segs.
    new_font_attr_segs = {}
    for line_index in self.font_attr_segs:
      new_font_attr_segs[other_num_lines + line_index] = (
          self.font_attr_segs[line_index])
    new_font_attr_segs.update(other.font_attr_segs)
    self._font_attr_segs = new_font_attr_segs

    # Merge the annotations.
    new_annotations = {}
    for key in self._annotations:
      if isinstance(key, int):
        new_annotations[other_num_lines + key] = (self.annotations[key])
      else:
        new_annotations[key] = other.annotations[key]

    new_annotations.update(other.annotations)
    self._annotations = new_annotations

  def append(self, line, font_attr_segs=None):
    """Append a single line of text.

    Args:
      line: (str) The text to be added to the end.
      font_attr_segs: (list of tuples) Font attribute segments of the appended
        line.
    """

    self._lines.append(line)
    if font_attr_segs:
      self._font_attr_segs[len(self._lines) - 1] = font_attr_segs

  def append_rich_line(self, rich_line):
    self.append(rich_line.text, rich_line.font_attr_segs)

  def prepend(self, line, font_attr_segs=None):
    """Prepend (i.e., add to the front) a single line of text.

    Args:
      line: (str) The text to be added to the front.
      font_attr_segs: (list of tuples) Font attribute segments of the appended
        line.
    """

    other = RichTextLines(line)
    if font_attr_segs:
      other.font_attr_segs[0] = font_attr_segs
    self._extend_before(other)

  def write_to_file(self, file_path):
    """Write the object itself to file, in a plain format.

    The font_attr_segs and annotations are ignored.

    Args:
      file_path: (str) path of the file to write to.
    """

    with gfile.Open(file_path, "w") as f:
      for line in self._lines:
        f.write(line + "\n")

  # TODO(cais): Add a method to allow appending to a line in RichTextLines with
  # both text and font_attr_segs.

def regex_find(orig_screen_output, regex, font_attr):
  """Perform regex match in rich text lines.

  Produces a new RichTextLines object with font_attr_segs containing highlighted
  regex matches.

  Example use cases include:
  1) search for specific items in a large list of items, and
  2) search for specific numerical values in a large tensor.

  Args:
    orig_screen_output: The original RichTextLines, in which the regex find
      is to be performed.
    regex: The regex used for matching.
    font_attr: Font attribute used for highlighting the found result.

  Returns:
    A modified copy of orig_screen_output.

  Raises:
    ValueError: If input str regex is not a valid regular expression.
  """
  new_screen_output = RichTextLines(
      orig_screen_output.lines,
      font_attr_segs=copy.deepcopy(orig_screen_output.font_attr_segs),
      annotations=orig_screen_output.annotations)

  try:
    re_prog = re.compile(regex)
  except re.error:
    raise ValueError("Invalid regular expression: \"%s\"" % regex)

  regex_match_lines = []
  for i, line in enumerate(new_screen_output.lines):
    find_it = re_prog.finditer(line)

    match_segs = []
    for match in find_it:
      match_segs.append((match.start(), match.end(), font_attr))

    if match_segs:
      if i not in new_screen_output.font_attr_segs:
        new_screen_output.font_attr_segs[i] = match_segs
      else:
        new_screen_output.font_attr_segs[i].extend(match_segs)
        new_screen_output.font_attr_segs[i] = sorted(
            new_screen_output.font_attr_segs[i], key=lambda x: x[0])
      regex_match_lines.append(i)

  new_screen_output.annotations[REGEX_MATCH_LINES_KEY] = regex_match_lines
  return new_screen_output

def wrap_rich_text_lines(inp, cols):
  """Wrap RichTextLines according to maximum number of columns.

  Produces a new RichTextLines object with the text lines, font_attr_segs and
  annotations properly wrapped. This ought to be used sparingly, as in most
  cases, command handlers producing RichTextLines outputs should know the
  screen/panel width via the screen_info kwarg and should produce properly
  length-limited lines in the output accordingly.

  Args:
    inp: Input RichTextLines object.
    cols: Number of columns, as an int.

  Returns:
    1) A new instance of RichTextLines, with line lengths limited to cols.
    2) A list of new (wrapped) line index. For example, if the original input
      consists of three lines and only the second line is wrapped, and it's
      wrapped into two lines, this return value will be: [0, 1, 3].
  Raises:
    ValueError: If inputs have invalid types.
  """

  new_line_indices = []

  if not isinstance(inp, RichTextLines):
    raise ValueError("Invalid type of input screen_output")

  if not isinstance(cols, int):
    raise ValueError("Invalid type of input cols")

  out = RichTextLines([])

  row_counter = 0  # Counter for new row index
  for i, line in enumerate(inp.lines):
    new_line_indices.append(out.num_lines())

    if i in inp.annotations:
      out.annotations[row_counter] = inp.annotations[i]

    if len(line) <= cols:
      # No wrapping.
      out.lines.append(line)
      if i in inp.font_attr_segs:
        out.font_attr_segs[row_counter] = inp.font_attr_segs[i]

      row_counter += 1
    else:
      # Wrap.
      wlines = []  # Wrapped lines.

      osegs = []
      if i in inp.font_attr_segs:
        osegs = inp.font_attr_segs[i]

      idx = 0
      while idx < len(line):
        if idx + cols > len(line):
          rlim = len(line)
        else:
          rlim = idx + cols

        wlines.append(line[idx:rlim])
        for seg in osegs:
          if (seg[0] < rlim) and (seg[1] >= idx):
            # Calculate left bound within wrapped line.
            if seg[0] >= idx:
              lb = seg[0] - idx
            else:
              lb = 0

            # Calculate right bound within wrapped line.
            if seg[1] < rlim:
              rb = seg[1] - idx
            else:
              rb = rlim - idx

            if rb > lb:  # Omit zero-length segments.
              wseg = (lb, rb, seg[2])
              if row_counter not in out.font_attr_segs:
                out.font_attr_segs[row_counter] = [wseg]
              else:
                out.font_attr_segs[row_counter].append(wseg)

        idx += cols
        row_counter += 1

      out.lines.extend(wlines)

  # Copy over keys of annotation that are not row indices.
  for key in inp.annotations:
    if not isinstance(key, int):
      out.annotations[key] = inp.annotations[key]

  return out, new_line_indices

class CommandHandlerRegistry:
  """Registry of command handlers for CLI.

  Handler methods (callables) for user commands can be registered with this
  class, which then is able to dispatch commands to the correct handlers and
  retrieve the RichTextLines output.

  For example, suppose you have the following handler defined:
    def echo(argv, screen_info=None):
      return RichTextLines(["arguments = %s" % " ".join(argv),
                            "screen_info = " + repr(screen_info)])

  you can register the handler with the command prefix "echo" and alias "e":
    registry = CommandHandlerRegistry()
    registry.register_command_handler("echo", echo,
        "Echo arguments, along with screen info", prefix_aliases=["e"])

  then to invoke this command handler with some arguments and screen_info, do:
    registry.dispatch_command("echo", ["foo", "bar"], screen_info={"cols": 80})

  or with the prefix alias:
    registry.dispatch_command("e", ["foo", "bar"], screen_info={"cols": 80})

  The call will return a RichTextLines object which can be rendered by a CLI.
  """

  HELP_COMMAND = "help"
  HELP_COMMAND_ALIASES = ["h"]
  VERSION_COMMAND = "version"
  VERSION_COMMAND_ALIASES = ["ver"]

  def __init__(self):
    # A dictionary from command prefix to handler.
    self._handlers = {}

    # A dictionary from prefix alias to prefix.
    self._alias_to_prefix = {}

    # A dictionary from prefix to aliases.
    self._prefix_to_aliases = {}

    # A dictionary from command prefix to help string.
    self._prefix_to_help = {}

    # Introductory text to help information.
    self._help_intro = None

    # Register a default handler for the command "help".
    self.register_command_handler(
        self.HELP_COMMAND,
        self._help_handler,
        "Print this help message.",
        prefix_aliases=self.HELP_COMMAND_ALIASES)

    # Register a default handler for the command "version".
    self.register_command_handler(
        self.VERSION_COMMAND,
        self._version_handler,
        "Print the versions of TensorFlow and its key dependencies.",
        prefix_aliases=self.VERSION_COMMAND_ALIASES)

  def register_command_handler(self,
                               prefix,
                               handler,
                               help_info,
                               prefix_aliases=None):
    """Register a callable as a command handler.

    Args:
      prefix: Command prefix, i.e., the first word in a command, e.g.,
        "print" as in "print tensor_1".
      handler: A callable of the following signature:
          foo_handler(argv, screen_info=None),
        where argv is the argument vector (excluding the command prefix) and
          screen_info is a dictionary containing information about the screen,
          such as number of columns, e.g., {"cols": 100}.
        The callable should return:
          1) a RichTextLines object representing the screen output.

        The callable can also raise an exception of the type CommandLineExit,
        which if caught by the command-line interface, will lead to its exit.
        The exception can optionally carry an exit token of arbitrary type.
      help_info: A help string.
      prefix_aliases: Aliases for the command prefix, as a list of str. E.g.,
        shorthands for the command prefix: ["p", "pr"]

    Raises:
      ValueError: If
        1) the prefix is empty, or
        2) handler is not callable, or
        3) a handler is already registered for the prefix, or
        4) elements in prefix_aliases clash with existing aliases.
        5) help_info is not a str.
    """

    if not prefix:
      raise ValueError("Empty command prefix")

    if prefix in self._handlers:
      raise ValueError(
          "A handler is already registered for command prefix \"%s\"" % prefix)

    # Make sure handler is callable.
    if not callable(handler):
      raise ValueError("handler is not callable")

    # Make sure that help info is a string.
    if not isinstance(help_info, str):
      raise ValueError("help_info is not a str")

    # Process prefix aliases.
    if prefix_aliases:
      for alias in prefix_aliases:
        if self._resolve_prefix(alias):
          raise ValueError(
              "The prefix alias \"%s\" clashes with existing prefixes or "
              "aliases." % alias)
        self._alias_to_prefix[alias] = prefix

      self._prefix_to_aliases[prefix] = prefix_aliases

    # Store handler.
    self._handlers[prefix] = handler

    # Store help info.
    self._prefix_to_help[prefix] = help_info

  def dispatch_command(self, prefix, argv, screen_info=None):
    """Handles a command by dispatching it to a registered command handler.

    Args:
      prefix: Command prefix, as a str, e.g., "print".
      argv: Command argument vector, excluding the command prefix, represented
        as a list of str, e.g.,
        ["tensor_1"]
      screen_info: A dictionary containing screen info, e.g., {"cols": 100}.

    Returns:
      An instance of RichTextLines or None. If any exception is caught during
      the invocation of the command handler, the RichTextLines will wrap the
      error type and message.

    Raises:
      ValueError: If
        1) prefix is empty, or
        2) no command handler is registered for the command prefix, or
        3) the handler is found for the prefix, but it fails to return a
          RichTextLines or raise any exception.
      CommandLineExit:
        If the command handler raises this type of exception, this method will
        simply pass it along.
    """
    if not prefix:
      raise ValueError("Prefix is empty")

    resolved_prefix = self._resolve_prefix(prefix)
    if not resolved_prefix:
      raise ValueError("No handler is registered for command prefix \"%s\"" %
                       prefix)

    handler = self._handlers[resolved_prefix]
    try:
      output = handler(argv, screen_info=screen_info)
    except CommandLineExit as e:
      raise e
    except SystemExit as e:
      # Special case for syntax errors caught by argparse.
      lines = ["Syntax error for command: %s" % prefix,
               "For help, do \"help %s\"" % prefix]
      output = RichTextLines(lines)

    except BaseException as e:  # pylint: disable=broad-except
      lines = ["Error occurred during handling of command: %s %s:" %
               (resolved_prefix, " ".join(argv)), "%s: %s" % (type(e), str(e))]

      # Include traceback of the exception.
      lines.append("")
      lines.extend(traceback.format_exc().split("\n"))

      output = RichTextLines(lines)

    if not isinstance(output, RichTextLines) and output is not None:
      raise ValueError(
          "Return value from command handler %s is not None or a RichTextLines "
          "instance" % str(handler))

    return output

  def is_registered(self, prefix):
    """Test if a command prefix or its alias is has a registered handler.

    Args:
      prefix: A prefix or its alias, as a str.

    Returns:
      True iff a handler is registered for prefix.
    """
    return self._resolve_prefix(prefix) is not None

  def get_help(self, cmd_prefix=None):
    """Compile help information into a RichTextLines object.

    Args:
      cmd_prefix: Optional command prefix. As the prefix itself or one of its
        aliases.

    Returns:
      A RichTextLines object containing the help information. If cmd_prefix
      is None, the return value will be the full command-line help. Otherwise,
      it will be the help information for the specified command.
    """
    if not cmd_prefix:
      # Print full help information, in sorted order of the command prefixes.
      help_info = RichTextLines([])
      if self._help_intro:
        # If help intro is available, show it at the beginning.
        help_info.extend(self._help_intro)

      sorted_prefixes = sorted(self._handlers)
      for cmd_prefix in sorted_prefixes:
        lines = self._get_help_for_command_prefix(cmd_prefix)
        lines.append("")
        lines.append("")
        help_info.extend(RichTextLines(lines))

      return help_info
    else:
      return RichTextLines(self._get_help_for_command_prefix(cmd_prefix))

  def set_help_intro(self, help_intro):
    """Set an introductory message to help output.

    Args:
      help_intro: (RichTextLines) Rich text lines appended to the
        beginning of the output of the command "help", as introductory
        information.
    """
    self._help_intro = help_intro

  def _help_handler(self, args, screen_info=None):
    """Command handler for "help".

    "help" is a common command that merits built-in support from this class.

    Args:
      args: Command line arguments to "help" (not including "help" itself).
      screen_info: (dict) Information regarding the screen, e.g., the screen
        width in characters: {"cols": 80}

    Returns:
      (RichTextLines) Screen text output.
    """

    _ = screen_info  # Unused currently.

    if not args:
      return self.get_help()
    elif len(args) == 1:
      return self.get_help(args[0])
    else:
      return RichTextLines(["ERROR: help takes only 0 or 1 input argument."])

  def _version_handler(self, args, screen_info=None):
    del args  # Unused currently.
    del screen_info  # Unused currently.
    return get_tensorflow_version_lines(include_dependency_versions=True)

  def _resolve_prefix(self, token):
    """Resolve command prefix from the prefix itself or its alias.

    Args:
      token: a str to be resolved.

    Returns:
      If resolvable, the resolved command prefix.
      If not resolvable, None.
    """
    if token in self._handlers:
      return token
    elif token in self._alias_to_prefix:
      return self._alias_to_prefix[token]
    else:
      return None

  def _get_help_for_command_prefix(self, cmd_prefix):
    """Compile the help information for a given command prefix.

    Args:
      cmd_prefix: Command prefix, as the prefix itself or one of its aliases.

    Returns:
      A list of str as the help information for cmd_prefix. If the cmd_prefix
        does not exist, the returned list of str will indicate that.
    """
    lines = []

    resolved_prefix = self._resolve_prefix(cmd_prefix)
    if not resolved_prefix:
      lines.append("Invalid command prefix: \"%s\"" % cmd_prefix)
      return lines

    lines.append(resolved_prefix)

    if resolved_prefix in self._prefix_to_aliases:
      lines.append(HELP_INDENT + "Aliases: " + ", ".join(
          self._prefix_to_aliases[resolved_prefix]))

    lines.append("")
    help_lines = self._prefix_to_help[resolved_prefix].split("\n")
    for line in help_lines:
      lines.append(HELP_INDENT + line)

    return lines

class TabCompletionRegistry:
  """Registry for tab completion responses."""

  def __init__(self):
    self._comp_dict = {}

  # TODO(cais): Rename method names with "comp" to "*completion*" to avoid
  # confusion.

  def register_tab_comp_context(self, context_words, comp_items):
    """Register a tab-completion context.

    Register that, for each word in context_words, the potential tab-completions
    are the words in comp_items.

    A context word is a pre-existing, completed word in the command line that
    determines how tab-completion works for another, incomplete word in the same
    command line.
    Completion items consist of potential candidates for the incomplete word.

    To give a general example, a context word can be "drink", and the completion
    items can be ["coffee", "tea", "water"]

    Note: A context word can be empty, in which case the context is for the
     top-level commands.

    Args:
      context_words: A list of context words belonging to the context being
        registered. It is a list of str, instead of a single string, to support
        synonym words triggering the same tab-completion context, e.g.,
        both "drink" and the short-hand "dr" can trigger the same context.
      comp_items: A list of completion items, as a list of str.

    Raises:
      TypeError: if the input arguments are not all of the correct types.
    """

    if not isinstance(context_words, list):
      raise TypeError("Incorrect type in context_list: Expected list, got %s" %
                      type(context_words))

    if not isinstance(comp_items, list):
      raise TypeError("Incorrect type in comp_items: Expected list, got %s" %
                      type(comp_items))

    # Sort the completion items on registration, so that later during
    # get_completions calls, no sorting will be necessary.
    sorted_comp_items = sorted(comp_items)

    for context_word in context_words:
      self._comp_dict[context_word] = sorted_comp_items

  def deregister_context(self, context_words):
    """Deregister a list of context words.

    Args:
      context_words: A list of context words to deregister, as a list of str.

    Raises:
      KeyError: if there are word(s) in context_words that do not correspond
        to any registered contexts.
    """

    for context_word in context_words:
      if context_word not in self._comp_dict:
        raise KeyError("Cannot deregister unregistered context word \"%s\"" %
                       context_word)

    for context_word in context_words:
      del self._comp_dict[context_word]

  def extend_comp_items(self, context_word, new_comp_items):
    """Add a list of completion items to a completion context.

    Args:
      context_word: A single completion word as a string. The extension will
        also apply to all other context words of the same context.
      new_comp_items: (list of str) New completion items to add.

    Raises:
      KeyError: if the context word has not been registered.
    """

    if context_word not in self._comp_dict:
      raise KeyError("Context word \"%s\" has not been registered" %
                     context_word)

    self._comp_dict[context_word].extend(new_comp_items)
    self._comp_dict[context_word] = sorted(self._comp_dict[context_word])

  def remove_comp_items(self, context_word, comp_items):
    """Remove a list of completion items from a completion context.

    Args:
      context_word: A single completion word as a string. The removal will
        also apply to all other context words of the same context.
      comp_items: Completion items to remove.

    Raises:
      KeyError: if the context word has not been registered.
    """

    if context_word not in self._comp_dict:
      raise KeyError("Context word \"%s\" has not been registered" %
                     context_word)

    for item in comp_items:
      self._comp_dict[context_word].remove(item)

  def get_completions(self, context_word, prefix):
    """Get the tab completions given a context word and a prefix.

    Args:
      context_word: The context word.
      prefix: The prefix of the incomplete word.

    Returns:
      (1) None if no registered context matches the context_word.
          A list of str for the matching completion items. Can be an empty list
          of a matching context exists, but no completion item matches the
          prefix.
      (2) Common prefix of all the words in the first return value. If the
          first return value is None, this return value will be None, too. If
          the first return value is not None, i.e., a list, this return value
          will be a str, which can be an empty str if there is no common
          prefix among the items of the list.
    """

    if context_word not in self._comp_dict:
      return None, None

    comp_items = self._comp_dict[context_word]
    comp_items = sorted(
        [item for item in comp_items if item.startswith(prefix)])

    return comp_items, self._common_prefix(comp_items)

  def _common_prefix(self, m):
    """Given a list of str, returns the longest common prefix.

    Args:
      m: (list of str) A list of strings.

    Returns:
      (str) The longest common prefix.
    """
    if not m:
      return ""

    s1 = min(m)
    s2 = max(m)
    for i, c in enumerate(s1):
      if c != s2[i]:
        return s1[:i]

    return s1

class CommandHistory:
  """Keeps command history and supports lookup."""

  _HISTORY_FILE_NAME = ".tfdbg_history"

  def __init__(self, limit=100, history_file_path=None):
    """CommandHistory constructor.

    Args:
      limit: Maximum number of the most recent commands that this instance
        keeps track of, as an int.
      history_file_path: (str) Manually specified path to history file. Used in
        testing.
    """

    self._commands = []
    self._limit = limit
    self._history_file_path = (
        history_file_path or self._get_default_history_file_path())
    self._load_history_from_file()

  def _load_history_from_file(self):
    if os.path.isfile(self._history_file_path):
      try:
        with open(self._history_file_path, "rt") as history_file:
          commands = history_file.readlines()
        self._commands = [command.strip() for command in commands
                          if command.strip()]

        # Limit the size of the history file.
        if len(self._commands) > self._limit:
          self._commands = self._commands[-self._limit:]
          with open(self._history_file_path, "wt") as history_file:
            for command in self._commands:
              history_file.write(command + "\n")
      except IOError:
        print("WARNING: writing history file failed.")

  def _add_command_to_history_file(self, command):
    try:
      with open(self._history_file_path, "at") as history_file:
        history_file.write(command + "\n")
    except IOError:
      pass

  @classmethod
  def _get_default_history_file_path(cls):
    return os.path.join(os.path.expanduser("~"), cls._HISTORY_FILE_NAME)

  def add_command(self, command):
    """Add a command to the command history.

    Args:
      command: The history command, as a str.

    Raises:
      TypeError: if command is not a str.
    """

    if self._commands and command == self._commands[-1]:
      # Ignore repeating commands in a row.
      return

    if not isinstance(command, str):
      raise TypeError("Attempt to enter non-str entry to command history")

    self._commands.append(command)

    if len(self._commands) > self._limit:
      self._commands = self._commands[-self._limit:]

    self._add_command_to_history_file(command)

  def most_recent_n(self, n):
    """Look up the n most recent commands.

    Args:
      n: Number of most recent commands to look up.

    Returns:
      A list of n most recent commands, or all available most recent commands,
      if n exceeds size of the command history, in chronological order.
    """

    return self._commands[-n:]

  def lookup_prefix(self, prefix, n):
    """Look up the n most recent commands that starts with prefix.

    Args:
      prefix: The prefix to lookup.
      n: Number of most recent commands to look up.

    Returns:
      A list of n most recent commands that have the specified prefix, or all
      available most recent commands that have the prefix, if n exceeds the
      number of history commands with the prefix.
    """

    commands = [cmd for cmd in self._commands if cmd.startswith(prefix)]

    return commands[-n:]

  # TODO(cais): Lookup by regex.

class MenuItem:
  """A class for an item in a text-based menu."""

  def __init__(self, caption, content, enabled=True):
    """Menu constructor.

    TODO(cais): Nested menu is currently not supported. Support it.

    Args:
      caption: (str) caption of the menu item.
      content: Content of the menu item. For a menu item that triggers
        a command, for example, content is the command string.
      enabled: (bool) whether this menu item is enabled.
    """

    self._caption = caption
    self._content = content
    self._enabled = enabled

  @property
  def caption(self):
    return self._caption

  @property
  def type(self):
    return self._node_type

  @property
  def content(self):
    return self._content

  def is_enabled(self):
    return self._enabled

  def disable(self):
    self._enabled = False

  def enable(self):
    self._enabled = True

class Menu:
  """A class for text-based menu."""

  def __init__(self, name=None):
    """Menu constructor.

    Args:
      name: (str or None) name of this menu.
    """

    self._name = name
    self._items = []

  def append(self, item):
    """Append an item to the Menu.

    Args:
      item: (MenuItem) the item to be appended.
    """
    self._items.append(item)

  def insert(self, index, item):
    self._items.insert(index, item)

  def num_items(self):
    return len(self._items)

  def captions(self):
    return [item.caption for item in self._items]

  def caption_to_item(self, caption):
    """Get a MenuItem from the caption.

    Args:
      caption: (str) The caption to look up.

    Returns:
      (MenuItem) The first-match menu item with the caption, if any.

    Raises:
      LookupError: If a menu item with the caption does not exist.
    """

    captions = self.captions()
    if caption not in captions:
      raise LookupError("There is no menu item with the caption \"%s\"" %
                        caption)

    return self._items[captions.index(caption)]

  def format_as_single_line(self,
                            prefix=None,
                            divider=" | ",
                            enabled_item_attrs=None,
                            disabled_item_attrs=None):
    """Format the menu as a single-line RichTextLines object.

    Args:
      prefix: (str) String added to the beginning of the line.
      divider: (str) The dividing string between the menu items.
      enabled_item_attrs: (list or str) Attributes applied to each enabled
        menu item, e.g., ["bold", "underline"].
      disabled_item_attrs: (list or str) Attributes applied to each
        disabled menu item, e.g., ["red"].

    Returns:
      (RichTextLines) A single-line output representing the menu, with
        font_attr_segs marking the individual menu items.
    """

    if (enabled_item_attrs is not None and
        not isinstance(enabled_item_attrs, list)):
      enabled_item_attrs = [enabled_item_attrs]

    if (disabled_item_attrs is not None and
        not isinstance(disabled_item_attrs, list)):
      disabled_item_attrs = [disabled_item_attrs]

    menu_line = prefix if prefix is not None else ""
    attr_segs = []

    for item in self._items:
      menu_line += item.caption
      item_name_begin = len(menu_line) - len(item.caption)

      if item.is_enabled():
        final_attrs = [item]
        if enabled_item_attrs:
          final_attrs.extend(enabled_item_attrs)
        attr_segs.append((item_name_begin, len(menu_line), final_attrs))
      else:
        if disabled_item_attrs:
          attr_segs.append(
              (item_name_begin, len(menu_line), disabled_item_attrs))

      menu_line += divider

    return RichTextLines(menu_line, font_attr_segs={0: attr_segs})

Output: rich_text_lines_from_rich_line_list

Output: regex_find

Output: get_tensorflow_version_lines

Output: wrap_rich_text_lines

Output: TabCompletionRegistry

Output: RichLine

Output: CommandHistory

Output: CommandHandlerRegistry

Output: Menu

Output: MenuItem

Output: CommandLineExit

Output: RichTextLines


Input:
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Reader class for tfdbg v2 debug events."""

import collections
import os
import threading

from tensorflow.core.protobuf import debug_event_pb2
from tensorflow.python.framework import errors
from tensorflow.python.framework import tensor_util
from tensorflow.python.lib.io import file_io
from tensorflow.python.lib.io import tf_record
from tensorflow.python.util import compat

DebugEventWithOffset = collections.namedtuple(
    "DebugEventWithOffset", "debug_event offset")

class DebugEventsReader:
  """Reader class for a tfdbg v2 DebugEvents directory."""

  # Number of digests after which a read lock is released and re-acquired during
  # serial reading of digests for SourceFiles, Execution, and
  # GraphExecutionTrace. This allows us to avoid releasing and re-acquiring the
  # lock too often (i.e., after each digest) and to minimize performance
  # penalty.
  _READER_RELEASE_PER = 100

  _METADATA_SUFFIX = ".metadata"
  _SOURCE_FILE_SUFFIX = ".source_files"
  _STACK_FRAMES_SUFFIX = ".stack_frames"
  _GRAPHS_SUFFIX = ".graphs"
  _EXECUTION_SUFFIX = ".execution"
  _GRAPH_EXECUTION_TRACES_SUFFIX = ".graph_execution_traces"

  def __init__(self, dump_root):
    if not file_io.is_directory(dump_root):
      raise ValueError("Specified dump_root is not a directory: %s" % dump_root)
    self._dump_root = dump_root
    self._metadata_paths = self._load_metadata_files()

    prefixes = [
        metadata_path[:-len(self._METADATA_SUFFIX)]
        for metadata_path in self._metadata_paths
    ]
    prefix = prefixes[0]  # This is the prefix of the main file set.
    self._source_files_path = compat.as_bytes(prefix + self._SOURCE_FILE_SUFFIX)
    self._stack_frames_path = compat.as_bytes(prefix +
                                              self._STACK_FRAMES_SUFFIX)
    self._graphs_path = compat.as_bytes(prefix + self._GRAPHS_SUFFIX)
    self._execution_path = compat.as_bytes(prefix + self._EXECUTION_SUFFIX)
    # There can be multiple .graph_execution_trace files each belonging
    # to a file set generated on an individual host, in the case of
    # a distributed TensorFlow job.
    # This is different from the other debug event files in the file set.
    self._graph_execution_traces_paths = [
        compat.as_bytes(prefix + self._GRAPH_EXECUTION_TRACES_SUFFIX)
        for prefix in prefixes
    ]
    self._readers = dict()  # A map from file path to reader.
    # A map from file path to current reading offset.
    self._reader_offsets = dict()
    # Lock for reader creation.
    self._readers_lock = threading.Lock()
    # Locks for read operation on individual readers.
    self._reader_read_locks = dict()

    self._offsets = dict()

  def _load_metadata_files(self):
    """Load and parse metadata files in the dump root.

    Check that all metadata files have a common tfdbg_run_id, and raise
    a ValueError if their tfdbg_run_ids differ.

    Returns:
      A list of metadata file paths in ascending order of their starting
        wall_time timestamp.
    """

    metadata_paths = file_io.get_matching_files(
        os.path.join(self._dump_root, "*%s" % self._METADATA_SUFFIX))
    if not metadata_paths:
      raise ValueError("Cannot find any tfdbg metadata file in directory: %s" %
                       self._dump_root)
    wall_times = []
    run_ids = []
    tensorflow_versions = []
    file_versions = []
    for metadata_path in metadata_paths:
      reader = tf_record.tf_record_random_reader(metadata_path)
      try:
        record = reader.read(0)[0]
        debug_event = debug_event_pb2.DebugEvent.FromString(record)
        wall_times.append(debug_event.wall_time)
        run_ids.append(debug_event.debug_metadata.tfdbg_run_id)
        tensorflow_versions.append(
            debug_event.debug_metadata.tensorflow_version
        )
        file_versions.append(debug_event.debug_metadata.file_version)
      except Exception as e:
        raise errors.DataLossError(
            None,
            None,
            "Error reading tfdbg metadata from paths %s" % metadata_paths,
        ) from e
      finally:
        reader.close()
    self._starting_wall_time = wall_times[0]
    self._tfdbg_run_id = run_ids[0]
    self._tensorflow_version = tensorflow_versions[0]
    self._file_version = file_versions[0]
    if len(metadata_paths) == 1:
      # Fast path for a common case (only one DebugEvent file set.)
      return metadata_paths

    num_no_id = len([run_id for run_id in run_ids if not run_id])
    if num_no_id:
      paths_without_run_id = [
          metadata_path
          for metadata_path, run_id in zip(metadata_paths, run_ids)
          if not run_id
      ]
      raise ValueError(
          "Found %d tfdbg metadata files and %d of them do not "
          "have tfdbg run ids. The metadata files without run ids are: %s" %
          (len(run_ids), num_no_id, paths_without_run_id))
    elif len(set(run_ids)) != 1:
      raise ValueError(
          "Unexpected: Found multiple (%d) tfdbg2 runs in directory %s" %
          (len(set(run_ids)), self._dump_root))
    # Return the metadata files in ascending order of their timestamps.
    paths_and_timestamps = sorted(
        zip(metadata_paths, wall_times), key=lambda t: t[1])
    self._starting_wall_time = paths_and_timestamps[0][1]
    return [path[0] for path in paths_and_timestamps]

  def starting_wall_time(self):
    """Get the starting timestamp of the instrumented TensorFlow program.

    When there are multiple hosts (i.e., multiple tfdbg file sets), the earliest
    timestamp among the file sets is returned. It is assumed to be the job that
    starts first (e.g., the coordinator).

    Returns:
      Starting timestamp in seconds since the epoch, as a float.
    """
    return self._starting_wall_time

  def tfdbg_run_id(self):
    """Get the run ID of the instrumented TensorFlow program."""
    return self._tfdbg_run_id

  def tensorflow_version(self):
    """Get the version string of TensorFlow that the debugged program ran on."""
    return self._tensorflow_version

  def tfdbg_file_version(self):
    """Get the tfdbg file format version."""
    return self._file_version

  def __enter__(self):
    return self

  def __exit__(self, exception_type, exception_value, traceback):
    del exception_type, exception_value, traceback  # Unused
    self.close()

  def _generic_iterator(self, file_path):
    """A helper method that makes an iterator given a debug-events file path.

    Repeated calls to this method create iterators that remember the last
    successful reading position (offset) for each given `file_path`. So the
    iterators are meant for incremental reading of the file.

    Args:
      file_path: Path to the file to create the iterator for.

    Yields:
      A tuple of (offset, debug_event_proto) on each `next()` call.
    """
    yield_count = 0
    reader = self._get_reader(file_path)
    read_lock = self._reader_read_locks[file_path]
    read_lock.acquire()
    try:
      while True:
        current_offset = self._reader_offsets[file_path]
        try:
          record, self._reader_offsets[file_path] = reader.read(current_offset)
        except (errors.DataLossError, IndexError):
          # We ignore partial read exceptions, because a record may be
          # truncated. The PyRandomRecordReader throws an `IndexError` when
          # offset goes out of bound.
          break
        yield DebugEventWithOffset(
            debug_event=debug_event_pb2.DebugEvent.FromString(record),
            offset=current_offset)
        yield_count += 1
        # The read lock must be periodically released to allow for concurrent
        # random reads. But we do so at a number of reads, instead of after
        # every single read, in order to minimize the performance penalty.
        if yield_count % self._READER_RELEASE_PER == 0:
          read_lock.release()
          read_lock.acquire()
    finally:
      read_lock.release()

  def _get_reader(self, file_path):
    """Get a random-access reader for TFRecords file at file_path."""
    file_path = compat.as_bytes(file_path)
    # The following code uses the double-checked locking pattern to optimize
    # the common case (where the reader is already initialized).
    if file_path not in self._readers:  # 1st check, without lock.
      with self._readers_lock:
        if file_path not in self._readers:  # 2nd check, with lock.
          self._readers[file_path] = tf_record.tf_record_random_reader(
              file_path)
          self._reader_read_locks[file_path] = threading.Lock()
          self._reader_offsets[file_path] = 0
    return self._readers[file_path]

  def source_files_iterator(self):
    return self._generic_iterator(self._source_files_path)

  def stack_frames_iterator(self):
    return self._generic_iterator(self._stack_frames_path)

  def graphs_iterator(self):
    return self._generic_iterator(self._graphs_path)

  def read_source_files_event(self, offset):
    """Read a DebugEvent proto at given offset from the .source_files file."""
    with self._reader_read_locks[self._source_files_path]:
      proto_string = self._get_reader(self._source_files_path).read(offset)[0]
    return debug_event_pb2.DebugEvent.FromString(proto_string)

  def read_graphs_event(self, offset):
    """Read a DebugEvent proto at a given offset from the .graphs file.

    Args:
      offset: Offset to read the DebugEvent proto from.

    Returns:
      A DebugEventProto.

    Raises:
      `errors.DataLossError` if offset is at a wrong location.
      `IndexError` if offset is out of range of the file.
    """
    return debug_event_pb2.DebugEvent.FromString(
        self._get_reader(self._graphs_path).read(offset)[0])

  def execution_iterator(self):
    return self._generic_iterator(self._execution_path)

  def read_execution_event(self, offset):
    """Read a DebugEvent proto at a given offset from the .execution file.

    Args:
      offset: Offset to read the DebugEvent proto from.

    Returns:
      A DebugEventProto.

    Raises:
      `errors.DataLossError` if offset is at a wrong location.
      `IndexError` if offset is out of range of the file.
    """
    with self._reader_read_locks[self._execution_path]:
      proto_string = self._get_reader(self._execution_path).read(offset)[0]
    return debug_event_pb2.DebugEvent.FromString(proto_string)

  def graph_execution_traces_iterators(self):
    return [
        self._generic_iterator(path)
        for path in self._graph_execution_traces_paths
    ]

  def read_graph_execution_traces_event(self, locator):
    """Read DebugEvent at given offset from given .graph_execution_traces file.

    Args:
      locator: A (file_index, offset) tuple that locates the DebugEvent
        containing the graph execution trace.

    Returns:
      A DebugEventProto.

    Raises:
      `errors.DataLossError` if offset is at a wrong location.
      `IndexError` if offset is out of range of the file.
    """
    file_index, offset = locator
    graph_execution_traces_path = self._graph_execution_traces_paths[file_index]
    with self._reader_read_locks[graph_execution_traces_path]:
      proto_string = self._get_reader(graph_execution_traces_path).read(
          offset)[0]
    return debug_event_pb2.DebugEvent.FromString(proto_string)

  def close(self):
    with self._readers_lock:
      file_paths = list(self._readers.keys())
      for file_path in file_paths:
        self._readers[file_path].close()
        del self._readers[file_path]

class BaseDigest:
  """Base class for digest.

  Properties:
    wall_time: A timestamp for the digest as a `float` (unit: s).
    locator: A datum that allows tracng the digest to its original
      location. It can be either of the two:
       1. Bytes offset from the beginning of the file as a single integer,
          for the case of all digests of the same kind coming from the same
          file.
       2. A tuple of a file index and a byte offset. This applies to case
          in which the same type of debugger data may come from multiple files,
          e.g., graph execution traces.
  """

  def __init__(self, wall_time, locator):
    self._wall_time = wall_time
    self._locator = locator

  @property
  def wall_time(self):
    return self._wall_time

  @property
  def locator(self):
    return self._locator

  def to_json(self):
    return {"wall_time": self.wall_time}

class ExecutionDigest(BaseDigest):
  """Light-weight digest summarizing top-level execution event.

  Use `DebugDataReader.read_execution(execution_digest)` to load the more
  detailed data object concerning the execution event (`Execution`).

  Properties:
    op_type: Type name of the executed op. In the case of the eager execution of
      an individual op, it is the name of the op (e.g., "MatMul").
      In the case of the execution of a tf.function (FuncGraph), this is the
      internally-generated name of the function (e.g.,
      "__inference_my_func_123").
    output_tensor_device_ids: IDs of the devices on which the output tensors of
      the execution reside. For no-output execution, this is `None`.
  """

  def __init__(self,
               wall_time,
               locator,
               op_type,
               output_tensor_device_ids=None):
    super().__init__(wall_time, locator)
    self._op_type = op_type
    self._output_tensor_device_ids = _tuple_or_none(output_tensor_device_ids)

  @property
  def op_type(self):
    return self._op_type

  @property
  def output_tensor_device_ids(self):
    return self._output_tensor_device_ids

  def to_json(self):
    output = super().to_json()
    output.update({
        "op_type": self.op_type,
        "output_tensor_device_ids": self.output_tensor_device_ids,
    })
    return output

def _tuple_or_none(data):
  return tuple(data) if data else None

class Execution(ExecutionDigest):
  """Detailed data relating to a top-level execution event.

  The execution is of an individual op or a tf.function, which may have any
  number of output tensors.

  Properties (beyond the base class `ExecutionDigest`):
    host_name: Name of the host on which the execution happened.
    stack_frame_ids: Reference IDs for stack frames, ordered from bottommost to
      topmost. Use `DebugDataReader.read_execution_stack_trace()` to load the
      detailed stack frames (filepath, lineno and function name).
    tensor_debug_mode: TensorDebugMode enum value, as an `int`.
    graph_id: ID of the executed FuncGraph (applicable only the execution of a
      tf.function). `None` for the eager execution of an individual op.
    input_tensor_ids: IDs of the input (eager) tensor(s) for this execution, if
      any. If the eager execution has no input tensor, this is `None`. Else,
      this is a `tuple` of `int`s.
    output_tensor_ids: IDs of the output (eager) tensor(s) from this execution,
      if any. If the eager execution produces no output tensor, this is `None`.
      Else, this is a `tuple` of `int`s.
    debug_tensor_values: Values of the debug tensor(s), applicable only to
      non-FULL_TENSOR tensor debug mode. A tuple of list of numbers. Each
      element of the tuple corresponds to an output tensor of the execution.
      See documentation of the various TensorDebugModes for the semantics of the
      numbers. If the eager execution produces no output tensor, this is
      `None`. Else, this is a `tuple` of `list`s.
  """

  def __init__(self,
               execution_digest,
               host_name,
               stack_frame_ids,
               tensor_debug_mode,
               graph_id=None,
               input_tensor_ids=None,
               output_tensor_ids=None,
               debug_tensor_values=None):
    super().__init__(
        execution_digest.wall_time,
        execution_digest.locator,
        execution_digest.op_type,
        output_tensor_device_ids=execution_digest.output_tensor_device_ids)
    self._host_name = host_name
    self._stack_frame_ids = tuple(stack_frame_ids)
    self._tensor_debug_mode = tensor_debug_mode
    self._graph_id = graph_id
    self._input_tensor_ids = _tuple_or_none(input_tensor_ids)
    self._output_tensor_ids = _tuple_or_none(output_tensor_ids)
    self._debug_tensor_values = _tuple_or_none(debug_tensor_values)

  @property
  def host_name(self):
    return self._host_name

  @property
  def stack_frame_ids(self):
    return self._stack_frame_ids

  @property
  def tensor_debug_mode(self):
    return self._tensor_debug_mode

  @property
  def graph_id(self):
    return self._graph_id

  @property
  def input_tensor_ids(self):
    return self._input_tensor_ids

  @property
  def num_outputs(self):
    return len(self._output_tensor_ids) if self._output_tensor_ids else 0

  @property
  def output_tensor_ids(self):
    return self._output_tensor_ids

  @property
  def debug_tensor_values(self):
    return self._debug_tensor_values

  def to_json(self):
    output = super().to_json()
    output.update({
        "host_name": self.host_name,
        "stack_frame_ids": self.stack_frame_ids,
        "tensor_debug_mode": self.tensor_debug_mode,
        "graph_id": self.graph_id,
        "input_tensor_ids": self.input_tensor_ids,
        "output_tensor_ids": self.output_tensor_ids,
        "debug_tensor_values": self.debug_tensor_values,
    })
    return output

class DebuggedGraph:
  """Data object representing debugging information about a tf.Graph.

  Includes `FuncGraph`s.

  Properties:
    name: Name of the graph (if any). May be `None` for non-function graphs.
    graph_id: Debugger-generated ID for the graph.
    inner_graph_ids: A list of the debugger-generated IDs for the graphs
      enclosed by this graph.
    outer_graph_id: If this graph is nested within an outer graph, ID of the
      outer graph. If this is an outermost graph, `None`.
  """

  def __init__(self,
               name,
               graph_id,
               outer_graph_id=None):
    self._name = name
    self._graph_id = graph_id
    self._outer_graph_id = outer_graph_id
    self._inner_graph_ids = []
    # A dictionary from op name to GraphOpCreationDigest.
    self._op_by_name = dict()
    # A dictionary mapping op to immediate downstream consumers.
    self._op_consumers = collections.defaultdict(list)

  def add_inner_graph_id(self, inner_graph_id):
    """Add the debugger-generated ID of a graph nested within this graph.

    Args:
      inner_graph_id: The debugger-generated ID of the nested inner graph.
    """
    assert isinstance(inner_graph_id, str)
    self._inner_graph_ids.append(inner_graph_id)

  def add_op(self, graph_op_creation_digest):
    """Add an op creation data object.

    Args:
      graph_op_creation_digest: A GraphOpCreationDigest data object describing
        the creation of an op inside this graph.
    """
    if graph_op_creation_digest.op_name in self._op_by_name:
      raise ValueError(
          "Duplicate op name: %s (op type: %s)" %
          (graph_op_creation_digest.op_name, graph_op_creation_digest.op_type))
    self._op_by_name[
        graph_op_creation_digest.op_name] = graph_op_creation_digest

  def add_op_consumer(self, src_op_name, src_slot, dst_op_name, dst_slot):
    """Add a consuming op for this op.

    Args:
      src_op_name: Name of the op of which the output tensor is being consumed.
      src_slot: 0-based output slot of the op being consumed.
      dst_op_name: Name of the consuming op (e.g., "Conv2D_3/BiasAdd")
      dst_slot: 0-based input slot of the consuming op that receives the tensor
        from this op.
    """
    self._op_consumers[src_op_name].append((src_slot, dst_op_name, dst_slot))

  @property
  def name(self):
    return self._name

  @property
  def graph_id(self):
    return self._graph_id

  @property
  def outer_graph_id(self):
    return self._outer_graph_id

  @property
  def inner_graph_ids(self):
    return self._inner_graph_ids

  def get_tensor_id(self, op_name, output_slot):
    """Get the ID of a symbolic tensor in this graph."""
    return self._op_by_name[op_name].output_tensor_ids[output_slot]

  def get_op_creation_digest(self, op_name):
    """Get the GraphOpCreationDigest for a op in the graph."""
    return self._op_by_name[op_name]

  def get_op_consumers(self, src_op_name):
    """Get all the downstream consumers of this op.

    Only data (non-control) edges are tracked.

    Args:
      src_op_name: Name of the op providing the tensor being consumed.

    Returns:
      A list of (src_slot, dst_op_name, dst_slot) tuples. In each item of
      the list:
        src_slot: 0-based output slot of the op of which the output tensor
          is being consumed.
        dst_op_name: Name of the consuming op (e.g., "Conv2D_3/BiasAdd")
        dst_slot: 0-based input slot of the consuming op that receives
          the tensor from this op.
    """
    return self._op_consumers[src_op_name]

  def to_json(self):
    return {
        "name": self.name,
        "graph_id": self.graph_id,
        "outer_graph_id": self._outer_graph_id,
        "inner_graph_ids": self._inner_graph_ids,
    }

class DebuggedDevice:
  """Debugger data regarding a device involved in the debugged program.

  Properties:
    device_name: Name of the device, as a str.
    device_id: An integer ID for the device, unique for each device within
      the scope of the debugged TensorFlow program.
  """

  def __init__(self,
               device_name,
               device_id):
    self._device_name = device_name
    self._device_id = device_id

  @property
  def device_name(self):
    return self._device_name

  @property
  def device_id(self):
    return self._device_id

  def to_json(self):
    return {
        "device_name": self._device_name,
        "device_id": self._device_id,
    }

class GraphOpCreationDigest(BaseDigest):
  """Data object describing the creation of an op inside a graph.

  For size efficiency, this digest object does not contain any stack frames or
  any references to them. To obtain the stack frames, use
  `DataReader.read_graph_op_creation_stack_trace()`.

  Properties (beyond the base class):
    graph_id: Debugger-generated ID of the immediately-enclosing graph.
    op_type: Type name of the op (e.g., "MatMul").
    op_name: Name of the op (e.g., "dense_1/MatMul").
    output_tensor_ids: Debugger-generated IDs for the output(s) of the op.
      If the op produces no output tensor, this is `None`. Else, this is a
      `tuple` of `int`s.
    input_names: Names of the input tensors to the op.
    device_name: The name of the device that the op is placed on (if available).
    host_name: Name of the host on which the op is created.
    stack_frame_ids: IDs of the frames of the stack trace at which the op
      is created.
  """

  def __init__(self,
               wall_time,
               locator,
               graph_id,
               op_type,
               op_name,
               output_tensor_ids,
               host_name,
               stack_frame_ids,
               input_names=None,
               device_name=None):
    super().__init__(wall_time, locator)
    self._graph_id = graph_id
    self._op_type = op_type
    self._op_name = op_name
    self._output_tensor_ids = _tuple_or_none(output_tensor_ids)
    self._host_name = host_name
    self._stack_frame_ids = stack_frame_ids
    self._input_names = _tuple_or_none(input_names)
    self._device_name = device_name

  @property
  def graph_id(self):
    return self._graph_id

  @property
  def op_type(self):
    return self._op_type

  @property
  def op_name(self):
    return self._op_name

  @property
  def output_tensor_ids(self):
    return self._output_tensor_ids

  @property
  def num_outputs(self):
    return len(self._output_tensor_ids) if self.output_tensor_ids else 0

  @property
  def input_names(self):
    return self._input_names

  @property
  def device_name(self):
    return self._device_name

  @property
  def host_name(self):
    return self._host_name

  @property
  def stack_frame_ids(self):
    return self._stack_frame_ids

  def to_json(self):
    output = super().to_json()
    output.update({
        "graph_id": self.graph_id,
        "op_type": self.op_type,
        "op_name": self.op_name,
        "output_tensor_ids": self.output_tensor_ids,
        "host_name": self.host_name,
        "stack_frame_ids": self.stack_frame_ids,
        "input_names": self.input_names,
        "device_name": self.device_name,
    })
    return output

class GraphExecutionTraceDigest(BaseDigest):
  """Light-weight summary of a intra-graph tensor execution event.

  Use `DebugDataReader.read_graph_execution_trace()` on this object to read more
  detailed data (`GraphExecutionTrace`).

  Properties (beyond the base class):
    op_type: Type name of the executed op (e.g., "Conv2D").
    op_name: Name of the op (e.g., "conv_2d_3/Conv2D").
    output_slot: Output slot index of the tensor.
    graph_id: The debugger-generated ID of the innermost (immediately-enclosing)
      graph.
  """

  def __init__(self, wall_time, locator, op_type, op_name, output_slot,
               graph_id):
    super().__init__(wall_time, locator)
    self._op_type = op_type
    self._op_name = op_name
    self._output_slot = output_slot
    self._graph_id = graph_id

  @property
  def op_type(self):
    return self._op_type

  @property
  def op_name(self):
    return self._op_name

  @property
  def output_slot(self):
    return self._output_slot

  @property
  def graph_id(self):
    return self._graph_id

  def to_json(self):
    output = super().to_json()
    output.update({
        "op_type": self.op_type,
        "op_name": self.op_name,
        "output_slot": self.output_slot,
        "graph_id": self.graph_id,
    })
    return output

class GraphExecutionTrace(GraphExecutionTraceDigest):
  """Detailed data object describing an intra-graph tensor execution.

  Attributes (in addition to GraphExecutionTraceDigest):
    graph_ids: The debugger-generated IDs of the graphs that enclose the
      executed op (tensor), ordered from the outermost to the innermost.
    graph_id: The debugger-generated ID of the innermost (immediately-enclosing)
      graph.
    tensor_debug_mode: TensorDebugMode enum value.
    debug_tensor_value: Debug tensor values (only for non-FULL_TENSOR
      tensor_debug_mode). A list of numbers. See the documentation of the
      TensorDebugModes for the semantics of the numbers.
    device_name: Device on which the tensor resides (if available)
  """

  def __init__(self,
               graph_execution_trace_digest,
               graph_ids,
               tensor_debug_mode,
               debug_tensor_value=None,
               device_name=None):
    super().__init__(graph_execution_trace_digest.wall_time,
                     graph_execution_trace_digest.locator,
                     graph_execution_trace_digest.op_type,
                     graph_execution_trace_digest.op_name,
                     graph_execution_trace_digest.output_slot,
                     graph_execution_trace_digest.graph_id)
    self._graph_ids = tuple(graph_ids)
    self._tensor_debug_mode = tensor_debug_mode
    self._debug_tensor_value = debug_tensor_value
    self._device_name = device_name

  @property
  def graph_ids(self):
    return self._graph_ids

  @property
  def graph_id(self):
    return self._graph_ids[-1]

  @property
  def tensor_debug_mode(self):
    return self._tensor_debug_mode

  @property
  def debug_tensor_value(self):
    return _tuple_or_none(self._debug_tensor_value)

  @property
  def device_name(self):
    return self._device_name

  def to_json(self):
    output = super().to_json()
    output.update({
        "graph_ids": self.graph_ids,
        "tensor_debug_mode": self.tensor_debug_mode,
        "debug_tensor_value": self.debug_tensor_value,
        "device_name": self.device_name,
    })
    return output

def _parse_tensor_value(tensor_proto, return_list=False):
  """Helper method for reading a tensor value from a tensor proto.

  The rationale for the distinction between `True` and `False value of
  `return_list` is as follows:
  - `return_list=True` is used for TensorDebugMode values other than
    FULL_TENSOR, e.g., CONCISE_HEALTH, SHAPE and FULL_HEATLH. Under
    those modes, the value is guaranteed (by contract) to be a 1D float64
    tensor.
  - `return_list=False` is used for the FULL_HEALTH TensorDebugMode
    specifically. Instead, we use `numpy.ndarray` to maximally preserve
    the shape, dtype and value information regarding the underlying tensor
    value. Under that mode, we don't use a python list to represent the
    tensor value because that can lead to loss of information (e.g., both
    float16 and float32 dtypes get mapped to Python floats).

  Args:
    tensor_proto: The TensorProto instance from which the tensor value will be
      loaded.
    return_list: Whether the return value will be a nested Python list that
      comes out from `numpy.ndarray.tolist()`.

  Returns:
    If parsing is successful, the tensor value as a `numpy.ndarray` or the
      nested Python list converted from it.
    If parsing fails, `None`.
  """
  try:
    ndarray = tensor_util.MakeNdarray(tensor_proto)
    return ndarray.tolist() if return_list else ndarray
  except TypeError:
    # Depending on tensor_debug_mode, certain dtype of tensors don't
    # have logged debug tensor values.
    return None

def _execution_digest_from_debug_event_proto(debug_event, locator):
  """Convert a DebugEvent proto into an ExecutionDigest data object."""
  return ExecutionDigest(
      debug_event.wall_time,
      locator,
      debug_event.execution.op_type,
      output_tensor_device_ids=(debug_event.execution.output_tensor_device_ids
                                or None))

def _execution_from_debug_event_proto(debug_event, locator):
  """Convert a DebugEvent proto into an Execution data object."""
  execution_proto = debug_event.execution

  debug_tensor_values = None
  if (execution_proto.tensor_debug_mode ==
      debug_event_pb2.TensorDebugMode.FULL_TENSOR):
    pass  # TODO(cais): Build tensor store.
  elif (execution_proto.tensor_debug_mode !=
        debug_event_pb2.TensorDebugMode.NO_TENSOR):
    debug_tensor_values = []
    for tensor_proto in execution_proto.tensor_protos:
      # TODO(cais): Refactor into a helper method.
      debug_tensor_values.append(
          _parse_tensor_value(tensor_proto, return_list=True))
  return Execution(
      _execution_digest_from_debug_event_proto(debug_event, locator),
      execution_proto.code_location.host_name,
      tuple(execution_proto.code_location.stack_frame_ids),
      execution_proto.tensor_debug_mode,
      graph_id=execution_proto.graph_id,
      input_tensor_ids=tuple(execution_proto.input_tensor_ids),
      output_tensor_ids=tuple(execution_proto.output_tensor_ids),
      debug_tensor_values=_tuple_or_none(debug_tensor_values))

class DebugDataReader:
  """A reader that reads structured debugging data in the tfdbg v2 format.

  The set of data read by an object of this class concerns the execution history
  of a tfdbg2-instrumented TensorFlow program.

  Note:
    - An object of this class incrementally reads data from files that belong to
      the tfdbg v2 DebugEvent file set. Calling `update()` triggers the reading
      from the last-successful reading positions in the files.
    - This object can be used as a context manager. Its `__exit__()` call
      closes the file readers cleanly.
  """

  def __init__(self, dump_root):
    self._reader = DebugEventsReader(dump_root)

    # TODO(cais): Implement pagination for memory constraints.
    self._execution_digests = []

    # Mapping (host_name, file_path) tuple to offset in the .source_files file.
    self._host_name_file_path_to_offset = collections.OrderedDict()
    # A dict mapping id to (host_name, file_path, lineno, func) tuple.
    self._stack_frame_by_id = dict()
    # Stores unprocessed stack frame IDs. This is necessary to handle the
    # case in which reading of the .stack_frames file gets ahead of the reading
    # of the .source_files file.
    self._unprocessed_stack_frames = dict()
    # A dict mapping id to DebuggedDevice objects.
    self._device_by_id = dict()
    # A dict mapping id to DebuggedGraph objects.
    self._graph_by_id = dict()
    self._graph_op_digests = []
    # TODO(cais): Implement pagination for memory constraints.
    self._graph_execution_trace_digests = []

    self._monitors = []

  def _add_monitor(self, monitor):
    self._monitors.append(monitor)

  def _load_source_files(self):
    """Incrementally read the .source_files DebugEvent file."""
    source_files_iter = self._reader.source_files_iterator()
    for debug_event, offset in source_files_iter:
      source_file = debug_event.source_file
      self._host_name_file_path_to_offset[
          (source_file.host_name, source_file.file_path)] = offset

  def _load_stack_frames(self):
    """Incrementally read the .stack_frames file.

    This must be called after _load_source_files().
    It assumes that the following contract is honored by the writer of the tfdbg
    v2 data file set:
      - Before a stack frame is written to the .stack_frames file, the
        corresponding source file information must have been written to the
        .source_files file first.
    """
    stack_frames_iter = self._reader.stack_frames_iterator()
    for debug_event, _ in stack_frames_iter:
      stack_frame_with_id = debug_event.stack_frame_with_id
      file_line_col = stack_frame_with_id.file_line_col
      self._unprocessed_stack_frames[stack_frame_with_id.id] = file_line_col
    # We do the processing in a separate stage, because the reading in the
    # .source_files file may sometimes get ahead of the .source_files file.
    unprocessed_stack_frame_ids = tuple(self._unprocessed_stack_frames.keys())
    for stack_frame_id in unprocessed_stack_frame_ids:
      file_line_col = self._unprocessed_stack_frames[stack_frame_id]
      if len(self._host_name_file_path_to_offset) > file_line_col.file_index:
        host_name, file_path = list(self._host_name_file_path_to_offset.keys())[
            file_line_col.file_index]
        self._stack_frame_by_id[stack_frame_id] = (
            host_name, file_path, file_line_col.line, file_line_col.func)
      del self._unprocessed_stack_frames[stack_frame_id]

  def _load_graphs(self):
    """Incrementally read the .graphs file.

    Compiles the DebuggedGraph and GraphOpCreation data.
    """
    graphs_iter = self._reader.graphs_iterator()
    for debug_event, offset in graphs_iter:
      if debug_event.graph_op_creation.ByteSize():
        op_creation_proto = debug_event.graph_op_creation
        op_digest = GraphOpCreationDigest(
            debug_event.wall_time,
            offset,
            op_creation_proto.graph_id,
            op_creation_proto.op_type,
            op_creation_proto.op_name,
            tuple(op_creation_proto.output_tensor_ids),
            op_creation_proto.code_location.host_name,
            tuple(op_creation_proto.code_location.stack_frame_ids),
            input_names=tuple(op_creation_proto.input_names))
        self._graph_op_digests.append(op_digest)
        debugged_graph = self._graph_by_id[op_creation_proto.graph_id]
        debugged_graph.add_op(op_digest)
        for dst_slot, input_name in enumerate(op_creation_proto.input_names):
          src_op_name, src_slot = input_name.split(":")
          debugged_graph.add_op_consumer(src_op_name, int(src_slot),
                                         op_creation_proto.op_name, dst_slot)

      elif debug_event.debugged_graph.ByteSize():
        graph_proto = debug_event.debugged_graph
        graph = DebuggedGraph(
            graph_proto.graph_name or None,
            graph_proto.graph_id,
            outer_graph_id=graph_proto.outer_context_id or None)
        self._graph_by_id[graph_proto.graph_id] = graph
        if graph_proto.outer_context_id:
          self._graph_by_id[
              graph_proto.outer_context_id].add_inner_graph_id(graph.graph_id)
      elif debug_event.debugged_device.ByteSize():
        device_proto = debug_event.debugged_device
        self._device_by_id[device_proto.device_id] = DebuggedDevice(
            device_proto.device_name, device_proto.device_id)

  def _load_graph_execution_traces(self):
    """Incrementally load the .graph_execution_traces file."""
    for i, traces_iter in enumerate(
        self._reader.graph_execution_traces_iterators()):
      for debug_event, offset in traces_iter:
        self._graph_execution_trace_digests.append(
            self._graph_execution_trace_digest_from_debug_event_proto(
                debug_event, (i, offset)))
        if self._monitors:
          graph_execution_trace = (
              self._graph_execution_trace_from_debug_event_proto(
                  debug_event, (i, offset)))
          for monitor in self._monitors:
            monitor.on_graph_execution_trace(
                len(self._graph_execution_trace_digests) - 1,
                graph_execution_trace)

  def _graph_execution_trace_digest_from_debug_event_proto(
      self, debug_event, locator):
    trace_proto = debug_event.graph_execution_trace
    op_name = trace_proto.op_name
    op_type = self._lookup_op_type(trace_proto.tfdbg_context_id, op_name)
    return GraphExecutionTraceDigest(
        debug_event.wall_time, locator, op_type, op_name,
        trace_proto.output_slot,
        debug_event.graph_execution_trace.tfdbg_context_id)

  def _graph_execution_trace_from_debug_event_proto(self, debug_event, locator):
    """Convert a DebugEvent proto into a GraphExecutionTrace data object."""
    trace_proto = debug_event.graph_execution_trace
    graph_ids = [trace_proto.tfdbg_context_id]
    # Walk up the chain of outer contexts (graphs), so as to include all of
    # their IDs
    while True:
      graph = self.graph_by_id(graph_ids[0])
      if graph.outer_graph_id:
        graph_ids.insert(0, graph.outer_graph_id)
      else:
        break

    if (trace_proto.tensor_debug_mode ==
        debug_event_pb2.TensorDebugMode.FULL_TENSOR):
      debug_tensor_value = None
    else:
      debug_tensor_value = _parse_tensor_value(
          trace_proto.tensor_proto, return_list=True)
    return GraphExecutionTrace(
        self._graph_execution_trace_digest_from_debug_event_proto(
            debug_event, locator),
        graph_ids=graph_ids,
        tensor_debug_mode=trace_proto.tensor_debug_mode,
        debug_tensor_value=debug_tensor_value,
        device_name=trace_proto.device_name or None)

  def _lookup_op_type(self, graph_id, op_name):
    """Lookup the type of an op by name and the immediately enclosing graph.

    Args:
      graph_id: Debugger-generated ID of the immediately-enclosing graph.
      op_name: Name of the op.

    Returns:
      Op type as a str.
    """
    return self._graph_by_id[graph_id].get_op_creation_digest(op_name).op_type

  def _load_execution(self):
    """Incrementally read the .execution file."""
    execution_iter = self._reader.execution_iterator()
    for debug_event, offset in execution_iter:
      self._execution_digests.append(
          _execution_digest_from_debug_event_proto(debug_event, offset))
      if self._monitors:
        execution = _execution_from_debug_event_proto(debug_event, offset)
        for monitor in self._monitors:
          monitor.on_execution(len(self._execution_digests) - 1, execution)

  def update(self):
    """Perform incremental read of the file set."""
    self._load_source_files()
    self._load_stack_frames()
    self._load_graphs()
    self._load_graph_execution_traces()
    self._load_execution()

  def source_file_list(self):
    """Get a list of source files known to the debugger data reader.

    Returns:
      A tuple of `(host_name, file_path)` tuples.
    """
    return tuple(self._host_name_file_path_to_offset.keys())

  def source_lines(self, host_name, file_path):
    """Read the line-by-line content of a source file.

    Args:
      host_name: Host name on which the source file is located.
      file_path: File path at which the source file is located.

    Returns:
      Lines of the source file as a `list` of `str`s.
    """
    offset = self._host_name_file_path_to_offset[(host_name, file_path)]
    return list(self._reader.read_source_files_event(offset).source_file.lines)

  def starting_wall_time(self):
    """Wall timestamp for when the debugged TensorFlow program started.

    Returns:
      Stating wall time as seconds since the epoch, as a `float`.
    """
    return self._reader.starting_wall_time()

  def tensorflow_version(self):
    """TensorFlow version used in the debugged TensorFlow program.

    Note: this is not necessarily the same as the version of TensorFlow used to
    load the DebugEvent file set.

    Returns:
      TensorFlow version used by the debugged program, as a `str`.
    """
    return self._reader.tensorflow_version()

  def tfdbg_run_id(self):
    """Get the debugger run ID of the debugged TensorFlow program."""
    return self._reader.tfdbg_run_id()

  def outermost_graphs(self):
    """Get the number of outer most graphs read so far."""
    return [graph for graph in self._graph_by_id.values()
            if not graph.outer_graph_id]

  def graph_by_id(self, graph_id):
    """Get a DebuggedGraph object by its ID."""
    return self._graph_by_id[graph_id]

  def device_name_by_id(self, device_id):
    """Get the name of a device by the debugger-generated ID of the device."""
    return self._device_by_id[device_id].device_name

  def device_name_map(self):
    """Get a map mapping device IDs to device names."""
    return {device_id: self._device_by_id[device_id].device_name
            for device_id in self._device_by_id}

  def graph_op_digests(self, op_type=None):
    """Get the list of the digests for graph-op creation so far.

    Args:
      op_type: Optional op type to filter the creation events with.

    Returns:
      A list of `GraphOpCreationDigest` objects.
    """
    if op_type is not None:
      return [digest for digest in self._graph_op_digests
              if digest.op_type == op_type]
    else:
      return self._graph_op_digests

  def graph_execution_traces(self, digest=False, begin=None, end=None):
    """Get all the intra-graph execution tensor traces read so far.

    Args:
      digest: Whether the results will be returned in the more light-weight
        digest form.
      begin: Optional beginning index for the requested traces or their digests.
        Python-style negative indices are supported.
      end: Optional ending index for the requested traces or their digests.
        Python-style negative indices are supported.

    Returns:
      If `digest`: a `list` of `GraphExecutionTraceDigest` objects.
      Else: a `list` of `GraphExecutionTrace` objects.
    """
    digests = self._graph_execution_trace_digests
    if begin is not None or end is not None:
      begin = begin or 0
      end = end or len(digests)
      digests = digests[begin:end]
    if digest:
      return digests
    else:
      return [self.read_graph_execution_trace(digest) for digest in digests]

  def num_graph_execution_traces(self):
    """Get the number of graph execution traces read so far."""
    return len(self._graph_execution_trace_digests)

  def executions(self, digest=False, begin=None, end=None):
    """Get `Execution`s or `ExecutionDigest`s this reader has read so far.

    Args:
      digest: Whether the results are returned in a digest form, i.e.,
        `ExecutionDigest` format, instead of the more detailed `Execution`
        format.
      begin: Optional beginning index for the requested execution data objects
        or their digests. Python-style negative indices are supported.
      end: Optional ending index for the requested execution data objects or
        their digests. Python-style negative indices are supported.

    Returns:
      If `digest`: a `list` of `ExecutionDigest` objects.
      Else: a `list` of `Execution` objects.
    """
    digests = self._execution_digests
    if begin is not None or end is not None:
      begin = begin or 0
      end = end or len(digests)
      digests = digests[begin:end]
    if digest:
      return digests
    else:
      # TODO(cais): Optimizer performance removing repeated file open/close.
      return [self.read_execution(digest) for digest in digests]

  def num_executions(self):
    """Get the number of execution events read so far."""
    return len(self._execution_digests)

  def read_execution(self, execution_digest):
    """Read a detailed Execution object."""
    debug_event = self._reader.read_execution_event(execution_digest.locator)
    return _execution_from_debug_event_proto(debug_event,
                                             execution_digest.locator)

  def read_graph_execution_trace(self, graph_execution_trace_digest):
    """Read the detailed graph execution trace.

    Args:
      graph_execution_trace_digest: A `GraphExecutionTraceDigest` object.

    Returns:
      The corresponding `GraphExecutionTrace` object.
    """
    debug_event = self._reader.read_graph_execution_traces_event(
        graph_execution_trace_digest.locator)
    return self._graph_execution_trace_from_debug_event_proto(
        debug_event, graph_execution_trace_digest.locator)

  def read_execution_stack_trace(self, execution):
    """Read the stack trace of a given Execution object.

    Args:
      execution: The Execution object of interest.

    Returns:
      1. The host name.
      2. The stack trace, as a list of (file_path, lineno, func) tuples.
    """
    host_name = self._stack_frame_by_id[execution.stack_frame_ids[0]][0]
    return (host_name, [
        self._stack_frame_by_id[frame_id][1:]
        for frame_id in execution.stack_frame_ids])

  def read_graph_op_creation_stack_trace(self, graph_op_creation_digest):
    """Read the stack trace of a given graph op creation object.

    Args:
      graph_op_creation_digest: The GraphOpCreationDigest object of interest.

    Returns:
      A tuple consisting of:
        1. The host name.
        2. The stack trace, as a list of (file_path, lineno, func) tuples.
    """
    return graph_op_creation_digest.host_name, [
        self._stack_frame_by_id[frame_id][1:]
        for frame_id in graph_op_creation_digest.stack_frame_ids
    ]

  # TODO(cais): Add graph_execution_digests() with an ExecutionDigest
  #   as a kwarg, to establish the association between top-level and intra-graph
  #   execution events.

  def execution_to_tensor_values(self, execution):
    """Read the full tensor values from an Execution or ExecutionDigest.

    Args:
      execution: An `ExecutionDigest` or `ExecutionDigest` object.

    Returns:
      A list of numpy arrays representing the output tensor values of the
        execution event.
    """
    debug_event = self._reader.read_execution_event(execution.locator)
    return [_parse_tensor_value(tensor_proto)
            for tensor_proto in debug_event.execution.tensor_protos]

  def graph_execution_trace_to_tensor_value(self, trace):
    """Read full tensor values from an Execution or ExecutionDigest.

    Args:
      trace: An `GraphExecutionTraceDigest` or `GraphExecutionTrace` object.

    Returns:
      A numpy array representing the output tensor value of the intra-graph
        tensor execution event.
    """
    debug_event = self._reader.read_graph_execution_traces_event(trace.locator)
    return _parse_tensor_value(debug_event.graph_execution_trace.tensor_proto)

  def symbolic_tensor_id(self, graph_id, op_name, output_slot):
    """Get the ID of a symbolic tensor.

    Args:
      graph_id: The ID of the immediately-enclosing graph.
      op_name: Name of the op.
      output_slot: Output slot as an int.

    Returns:
      The ID of the symbolic tensor as an int.
    """
    return self._graph_by_id[graph_id].get_tensor_id(op_name, output_slot)

  def graph_execution_trace_to_tensor_id(self, trace):
    """Get symbolic tensor ID from a GraphExecutoinTraceDigest object."""
    return self.symbolic_tensor_id(
        trace.graph_id, trace.op_name, trace.output_slot)

  def __enter__(self):
    return self

  def __exit__(self, exception_type, exception_value, traceback):
    del exception_type, exception_value, traceback  # Unused
    self._reader.close()

Output: _execution_digest_from_debug_event_proto

Output: _parse_tensor_value

Output: _execution_from_debug_event_proto

Output: _tuple_or_none

Output: ExecutionDigest

Output: GraphExecutionTraceDigest

Output: GraphOpCreationDigest

Output: DebuggedGraph

Output: BaseDigest

Output: GraphExecutionTrace

Output: DebugDataReader

Output: Execution

Output: DebugEventsReader

Output: DebuggedDevice