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.
# ==============================================================================
"""Framework of debug wrapper sessions.

A debug wrapper session is a wrapper around a TensorFlow Python Session.
The wrapper preserves the Session interface, most importantly the run() method,
while providing abilities to:
a) Intercept a run() call to a wrapped session and insert debug tensor watches
   according to externally-specified debug URLs.

b) Release control to an external (i.e., non-Session) object before and after
   the run() call, so that the external object can perform actions such as
   launching a UI to let users inspect the intermediate tensors and partition
   graphs from the run() call.

c) (To be implemented in a future CL) Enter an instruction loop to let an
   external object (e.g., remote client) launch run() and cont() calls
   remotely.

*** The lifetime of a debug wrapper session: ***

1) The wrapper session is created by calling the constructor with a
   wrapped (normal) session as the argument:
     wrapper = FooDebugWrapperSession(sess)
   wherein FooDebugWrapperSession is a concrete subclass implementing the
   abstract BaseDebugWrapperSession class below.

2) Near the end of the constructor call, the on_session_init() callback is
   invoked, with a OnSessionInitRequest object as the argument. The object
   carries the wrapped (normal) session object.

3) The callback handles the request and returns a OnSessionInitResponse
   object with an action field, directing the wrapper session what to do next.

If the action field in the OnSessionInitResponse is PROCEED, the constructor
returns. Control is released back to the caller of the constructor, which can
invoke run() method of wrapper session with the same syntax as a non-wrapped
session, e.g.,:
  wrapper.run(fetches, feed_dict=feeds, options=run_options)

Below, A1 - A2 is the lifetime of a wrapper run() call if the action is
PROCEED:

A1) Right at the start of each run() call, the on_run_start() callback is
    invoked, with an OnRunStartRequest object carrying information such as
    the fetches, the feed dict, the run options and run metadata used in
    this run call, along with a count of how many run calls has occurred
    on this wrapper session. The callback then returns an OnRunStartResponse
    object, of which the action field directs what the wrapper session
    actually will do of the run() call.

    If the action is DEBUG_RUN, a debugged (tensor-watched) run will ensue,
    with the debug URLs supplied in the debug_urls field of the response.
    These can be file:// or grpc:// URLs, for example.

    If the action is NON_DEBUG_RUN, a non-debug (normal) run will ensue.

A2) Right before the run() returns, the on_run_end() callback is invoked,
    with an OnRunEndRequest object as the argument, which carries information
    including the actual action performed in the wrapper run() call and the
    run_metadata from the run() call.

However, if the action field in OnSessionInitResponse is
REMOTE_INSTR_LOOP, the constructor will automatically invoke an instruction loop
that gives the control to a remote caller.

In the remote instruction loop, the following steps will happen:

B1) Callback on_instr_start() is invoked. The callback will return an
    OnInstrStartResponse object with an action field which can order one of
    the following actions:
        i) a run() call with fetches, feeds and debug_urls specified.
       ii) exit the instruction loop.

B2) The wrapper session carries out the action specified above.

B3) If still in the instruction loop, the wrapper session invokes the
    on_instr_end() callback. After the on_instr_end() callback returns, jump
    back to B1.

TODO(cais): Implemented the instruction loop in B1 - B3.

"""

import abc
import re
import threading

from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import session
from tensorflow.python.debug.lib import debug_utils
from tensorflow.python.framework import errors
from tensorflow.python.framework import stack
from tensorflow.python.platform import tf_logging
from tensorflow.python.training import monitored_session
from tensorflow.python.util import nest
from tensorflow.python.util.compat import collections_abc

# Helper function.
def _check_type(obj, expected_types):
  """Check if an object is of the expected type.

  Args:
    obj: The object being checked.
    expected_types: (`type` or an iterable of `type`s) The expected `type`(s)
      of obj.

  Raises:
      TypeError: If obj is not an instance of expected_type.
  """
  if not isinstance(obj, expected_types):
    raise TypeError("Expected type %s; got type %s" %
                    (expected_types, type(obj)))

class OnSessionInitRequest:
  """Request to an on-session-init callback.

  This callback is invoked during the __init__ call to a debug-wrapper session.
  """

  def __init__(self, sess):
    """Constructor.

    Args:
      sess: A tensorflow Session object.
    """

    _check_type(sess, (session.BaseSession, monitored_session.MonitoredSession))
    self.session = sess

class OnSessionInitAction:
  """Enum-like values for possible action to take on session init."""

  # Proceed, without special actions, in the wrapper session initialization.
  # What action the wrapper session performs next is determined by the caller
  # of the wrapper session. E.g., it can call run().
  PROCEED = "proceed"

  # Instead of letting the caller of the wrapper session determine what actions
  # the wrapper session will perform next, enter a loop to receive instructions
  # from a remote client.
  # For example, TensorBoard visual debugger can use this action so that it can
  # launch session.run() calls remotely.
  REMOTE_INSTR_LOOP = "remote_instr_loop"

class OnSessionInitResponse:
  """Response from an on-session-init callback."""

  def __init__(self, action):
    """Constructor.

    Args:
      action: (`OnSessionInitAction`) Debugger action to take on session init.
    """
    _check_type(action, str)
    self.action = action

class OnRunStartRequest:
  """Request to an on-run-start callback.

  This callback is invoked during a run() call of the debug-wrapper
  session, immediately after the run() call counter is incremented.
  """

  def __init__(self, fetches, feed_dict, run_options, run_metadata,
               run_call_count, is_callable_runner=False):
    """Constructor of `OnRunStartRequest`.

    Args:
      fetches: Fetch targets of the run() call.
      feed_dict: The feed dictionary to the run() call.
      run_options: RunOptions input to the run() call.
      run_metadata: RunMetadata input to the run() call.
        The above four arguments are identical to the input arguments to the
        run() method of a non-wrapped TensorFlow session.
      run_call_count: 1-based count of how many run calls (including this one)
        has been invoked.
      is_callable_runner: (bool) whether a runner returned by
        Session.make_callable is being run.
    """
    self.fetches = fetches
    self.feed_dict = feed_dict
    self.run_options = run_options
    self.run_metadata = run_metadata
    self.run_call_count = run_call_count
    self.is_callable_runner = is_callable_runner

class OnRunStartAction:
  """Enum-like values for possible action to take on start of a run() call."""

  # Run once with debug tensor-watching.
  DEBUG_RUN = "debug_run"

  # Run once with profiler.
  PROFILE_RUN = "profile_run"

  # Run without debug tensor-watching.
  NON_DEBUG_RUN = "non_debug_run"

class OnRunStartResponse:
  """Request from an on-run-start callback.

  The caller of the callback can use this response object to specify what
  action the debug-wrapper session actually takes on the run() call.
  """

  def __init__(self,
               action,
               debug_urls,
               debug_ops="DebugIdentity",
               node_name_regex_allowlist=None,
               op_type_regex_allowlist=None,
               tensor_dtype_regex_allowlist=None,
               tolerate_debug_op_creation_failures=False):
    """Constructor of `OnRunStartResponse`.

    Args:
      action: (`OnRunStartAction`) the action actually taken by the wrapped
        session for the run() call.
      debug_urls: (`list` of `str`) debug_urls used in watching the tensors
        during the run() call.
      debug_ops: (`str` or `list` of `str`) Debug op(s) to be used by the
        debugger.
      node_name_regex_allowlist: Regular-expression allowlist for node
        name.
      op_type_regex_allowlist: Regular-expression allowlist for op type.
      tensor_dtype_regex_allowlist: Regular-expression allowlist for tensor
        dtype.
      tolerate_debug_op_creation_failures: Whether debug op creation failures
        are to be tolerated.
    """

    _check_type(action, str)
    self.action = action

    _check_type(debug_urls, list)
    self.debug_urls = debug_urls

    self.debug_ops = debug_ops

    self.node_name_regex_allowlist = node_name_regex_allowlist
    self.op_type_regex_allowlist = op_type_regex_allowlist
    self.tensor_dtype_regex_allowlist = tensor_dtype_regex_allowlist
    self.tolerate_debug_op_creation_failures = (
        tolerate_debug_op_creation_failures)

class OnRunEndRequest:
  """Request to an on-run-end callback.

  The callback is invoked immediately before the wrapped run() call ends.
  """

  def __init__(self,
               performed_action,
               run_metadata=None,
               client_graph_def=None,
               tf_error=None):
    """Constructor for `OnRunEndRequest`.

    Args:
      performed_action: (`OnRunStartAction`) Actually-performed action by the
        debug-wrapper session.
      run_metadata: run_metadata output from the run() call (if any).
      client_graph_def: (GraphDef) GraphDef from the client side, i.e., from
        the python front end of TensorFlow. Can be obtained with
        session.graph.as_graph_def().
      tf_error: (errors.OpError subtypes) TensorFlow OpError that occurred
        during the run (if any).
    """

    _check_type(performed_action, str)
    self.performed_action = performed_action

    if run_metadata is not None:
      _check_type(run_metadata, config_pb2.RunMetadata)
    self.run_metadata = run_metadata
    self.client_graph_def = client_graph_def
    self.tf_error = tf_error

class OnRunEndResponse:
  """Response from an on-run-end callback."""

  def __init__(self):

    # Currently only a placeholder.
    pass

class BaseDebugWrapperSession(session.SessionInterface, metaclass=abc.ABCMeta):
  """Base class of debug-wrapper session classes.

  Concrete classes that inherit from this class need to implement the abstract
  methods such as on_session_init, on_run_start and on_run_end.
  """

  def __init__(self, sess, thread_name_filter=None,
               pass_through_operrors=False):
    """Constructor of `BaseDebugWrapperSession`.

    Args:
      sess: An (unwrapped) TensorFlow session instance. It should be a subtype
        of `BaseSession` or `tf.MonitoredSession`.
      thread_name_filter: Regular-expression filter (allowlist) for name(s) of
        thread(s) on which the wrapper session will be active. This regular
        expression is used in a start-anchored fashion on the thread name, i.e.,
        by applying the `match` method of the compiled pattern. The default
        `None` means that the wrapper session will be active on all threads.
        E.g., r"MainThread$", r"QueueRunnerThread.*".
      pass_through_operrors: If True, all captured OpErrors will be
        propagated.  By default this captures all OpErrors.

    Raises:
      ValueError: On invalid `OnSessionInitAction` value.
      NotImplementedError: If a non-DirectSession sess object is received.
    """

    _check_type(sess, (session.BaseSession, monitored_session.MonitoredSession))

    # The session being wrapped.
    self._sess = sess
    self._thread_name_filter_pattern = (re.compile(thread_name_filter)
                                        if thread_name_filter else None)
    # TODO(cais/kstevens): Unittest this pass through feature.
    self._pass_through_operrors = pass_through_operrors

    # Keeps track of number of run calls that have been performed on this
    # debug-wrapper session. The count can be used for purposes such as
    # displaying the state of the Session in a UI and determining a run
    # number-dependent debug URL.
    self._run_call_count = 0

    # Invoke on-session-init callback.
    response = self.on_session_init(OnSessionInitRequest(self._sess))
    _check_type(response, OnSessionInitResponse)

    if response.action == OnSessionInitAction.PROCEED:
      pass
    elif response.action == OnSessionInitAction.REMOTE_INSTR_LOOP:
      # TODO(cais): Implement REMOTE_INSTR_LOOP
      raise NotImplementedError(
          "OnSessionInitAction REMOTE_INSTR_LOOP has not been "
          "implemented.")
    else:
      raise ValueError(
          "Invalid OnSessionInitAction value: %s" % response.action)

    self._default_session_context_manager = None

    # A cache for callables created from CallableOptions.
    self._cached_callables_from_options = {}

  @property
  def graph(self):
    return self._sess.graph

  @property
  def graph_def(self):
    return self._sess.graph_def

  @property
  def sess_str(self):
    return self._sess.sess_str

  @property
  def session(self):
    return self._sess

  def run(self,
          fetches,
          feed_dict=None,
          options=None,
          run_metadata=None,
          callable_runner=None,
          callable_runner_args=None,
          callable_options=None):
    """Wrapper around Session.run() that inserts tensor watch options.

    Args:
      fetches: Same as the `fetches` arg to regular `Session.run()`.
      feed_dict: Same as the `feed_dict` arg to regular `Session.run()`.
      options: Same as the `options` arg to regular `Session.run()`.
      run_metadata: Same as the `run_metadata` arg to regular `Session.run()`.
      callable_runner: A `callable` returned by `Session.make_callable()`.
        If not `None`, `fetches` and `feed_dict` must both be `None`.
        Mutually exclusive with `callable_options`.
      callable_runner_args: An optional list of arguments to `callable_runner`
        or for `callable_options`.
      callable_options: An instance of `config_pb2.CallableOptions`, to be
        used with `Session._make_callable_from_options()`. Mutually exclusive
        with `callable_runner`.

    Returns:
      Simply forwards the output of the wrapped `Session.run()` call.

    Raises:
      ValueError: On invalid `OnRunStartAction` value. Or if `callable_runner`
        is not `None` and either or both of `fetches` and `feed_dict` is `None`.
    """
    if callable_runner and callable_options:
      raise ValueError(
          "callable_runner and callable_options are mutually exclusive, but "
          "are both specified in this call to BaseDebugWrapperSession.run().")

    if callable_runner and (fetches or feed_dict):
      raise ValueError(
          "callable_runner and fetches/feed_dict are mutually exclusive, "
          "but are used simultaneously.")
    elif callable_options and (fetches or feed_dict):
      raise ValueError(
          "callable_options and fetches/feed_dict are mutually exclusive, "
          "but are used simultaneously.")

    self.increment_run_call_count()

    def is_empty(x):
      """Check whether a possibly nested structure is empty."""
      if not nest.is_nested(x):
        return False
      if isinstance(x, collections_abc.Mapping):
        return is_empty(list(x.values()))
      for item in x:
        if not is_empty(item):
          return False
      return True

    empty_fetches = is_empty(fetches)
    if empty_fetches:
      tf_logging.info(
          "Due to empty fetches, tfdbg Session wrapper is letting a "
          "Session.run pass through without any debugging actions.")
    if self._is_disabled_thread() or empty_fetches:
      if callable_runner:
        return callable_runner(*callable_runner_args)
      elif callable_options:
        # pylint:disable=protected-access
        return self._sess._make_callable_from_options(
            callable_options)(*callable_runner_args)
        # pylint:enable=protected-access
      else:
        return self._sess.run(fetches,
                              feed_dict=feed_dict,
                              options=options,
                              run_metadata=run_metadata)

    # Invoke on-run-start callback and obtain response.
    run_start_resp = self.on_run_start(
        OnRunStartRequest(fetches, feed_dict, options, run_metadata,
                          self._run_call_count,
                          is_callable_runner=bool(callable_runner)))
    _check_type(run_start_resp, OnRunStartResponse)

    if run_start_resp.action == OnRunStartAction.DEBUG_RUN:
      retvals, run_end_req = self._run_with_debugging(
          run_start_resp, fetches, feed_dict, options, run_metadata,
          callable_runner, callable_runner_args, callable_options)
    elif run_start_resp.action == OnRunStartAction.PROFILE_RUN:
      retvals, run_end_req = self._run_with_profiling(
          run_start_resp, fetches, feed_dict, options, run_metadata,
          callable_runner, callable_runner_args, callable_options)
    elif run_start_resp.action == OnRunStartAction.NON_DEBUG_RUN:
      # Invoke run() method of the wrapped session.
      if callable_runner:
        retvals = callable_runner(*callable_runner_args)
      elif callable_options:
        # pylint:disable=protected-access
        callable_object = self._sess._make_callable_from_options(
            callable_options)
        # pylint:enable=protected-access
        retvals = callable_object(*callable_runner_args)
      else:
        retvals = self._sess.run(
            fetches,
            feed_dict=feed_dict,
            options=options,
            run_metadata=run_metadata)

      # Prepare arg for the on-run-end callback.
      run_end_req = OnRunEndRequest(run_start_resp.action)
    else:
      raise ValueError(
          "Invalid OnRunStartAction value: %s" % run_start_resp.action)

    # Invoke on-run-end callback and obtain response.
    run_end_resp = self.on_run_end(run_end_req)
    _check_type(run_end_resp, OnRunEndResponse)
    # Currently run_end_resp is only a placeholder. No action is taken on it.

    return retvals

  def _run_with_debugging(self,
                          run_start_resp,
                          fetches,
                          feed_dict,
                          options,
                          run_metadata,
                          callable_runner,
                          callable_runner_args,
                          callable_options):
    """Perform a session.run() or callable with debugging."""
    # Decorate RunOption to fill in debugger tensor watch specifications.
    decorated_run_options = None
    if callable_options:
      callable_options_id = id(callable_options)
      if callable_options_id not in self._cached_callables_from_options:
        # Make a copy of callable_options to avoid mutating it.
        new_callable_options = config_pb2.CallableOptions()
        new_callable_options.CopyFrom(callable_options)
        decorated_run_options = new_callable_options.run_options
    else:
      decorated_run_options = options or config_pb2.RunOptions()

    run_metadata = run_metadata or config_pb2.RunMetadata()

    if decorated_run_options:
      self._decorate_run_options_for_debug(
          decorated_run_options,
          run_start_resp.debug_urls,
          debug_ops=run_start_resp.debug_ops,
          node_name_regex_allowlist=(run_start_resp.node_name_regex_allowlist),
          op_type_regex_allowlist=run_start_resp.op_type_regex_allowlist,
          tensor_dtype_regex_allowlist=(
              run_start_resp.tensor_dtype_regex_allowlist),
          tolerate_debug_op_creation_failures=(
              run_start_resp.tolerate_debug_op_creation_failures))

    # Invoke the run() method of the wrapped Session. Catch any TensorFlow
    # runtime errors.
    tf_error = None
    try:
      if callable_runner:
        retvals = callable_runner(*callable_runner_args,
                                  options=decorated_run_options,
                                  run_metadata=run_metadata)
      elif callable_options:
        # pylint:disable=protected-access
        if callable_options_id in self._cached_callables_from_options:
          callable_object = self._cached_callables_from_options[
              callable_options_id]
        else:
          callable_object = self._sess._make_callable_from_options(
              new_callable_options)
          self._cached_callables_from_options[
              callable_options_id] = callable_object
        # pylint:enable=protected-access
        retvals = callable_object(
            *callable_runner_args, run_metadata=run_metadata)
      else:
        retvals = self._sess.run(fetches,
                                 feed_dict=feed_dict,
                                 options=decorated_run_options,
                                 run_metadata=run_metadata)
    except errors.OpError as op_error:
      if self._pass_through_operrors:
        raise op_error
      tf_error = op_error
      retvals = op_error

    return retvals, OnRunEndRequest(
        run_start_resp.action,
        run_metadata=run_metadata,
        client_graph_def=self._sess.graph.as_graph_def(),
        tf_error=tf_error)

  def _run_with_profiling(self,
                          run_start_resp,
                          fetches,
                          feed_dict,
                          options,
                          run_metadata,
                          callable_runner,
                          callable_runner_args,
                          callable_options):
    """Perform a session.run() or callable with profiling."""
    # Decorate RunOption to fill in debugger tensor watch specifications.
    decorated_run_options = None
    if callable_options:
      callable_options_id = id(callable_options)
      if callable_options_id not in self._cached_callables_from_options:
        # Make a copy of callable_options to avoid mutating it.
        new_callable_options = config_pb2.CallableOptions()
        new_callable_options.CopyFrom(callable_options)
        decorated_run_options = new_callable_options.run_options
    else:
      decorated_run_options = options or config_pb2.RunOptions()
    self._decorate_run_options_for_profile(decorated_run_options)

    run_metadata = run_metadata or config_pb2.RunMetadata()
    if callable_runner:
      retvals = callable_runner(*callable_runner_args,
                                options=decorated_run_options,
                                run_metadata=run_metadata)
    elif callable_options:
      # pylint:disable=protected-access
      callable_object = self._sess._make_callable_from_options(
          new_callable_options)
      # pylint:enable=protected-access
      retvals = callable_object(
          *callable_runner_args, run_metadata=run_metadata)
    else:
      retvals = self._sess.run(fetches,
                               feed_dict=feed_dict,
                               options=decorated_run_options,
                               run_metadata=run_metadata)
    return retvals, OnRunEndRequest(
        run_start_resp.action,
        run_metadata=run_metadata,
        client_graph_def=self._sess.graph.as_graph_def())

  def _is_disabled_thread(self):
    thread_name = threading.current_thread().name or ""
    return (self._thread_name_filter_pattern and
            not self._thread_name_filter_pattern.match(thread_name))

  def run_step_fn(self, step_fn):
    return step_fn(
        monitored_session.MonitoredSession.StepContext(self._sess, self.run))

  def partial_run_setup(self, fetches, feeds=None):
    """Sets up the feeds and fetches for partial runs in the session."""
    raise NotImplementedError(
        "partial_run_setup is not implemented for debug-wrapper sessions.")

  def partial_run(self, handle, fetches, feed_dict=None):
    raise NotImplementedError(
        "partial_run is not implemented for debug-wrapper sessions.")

  def list_devices(self, *args, **kwargs):
    return self._sess.list_devices(*args, **kwargs)

  def reset(self, *args, **kwargs):
    return self._sess.reset(*args, **kwargs)

  def make_callable(self,
                    fetches,
                    feed_list=None,
                    accept_options=False):
    runner = self._sess.make_callable(
        fetches, feed_list=feed_list, accept_options=True)
    def wrapped_runner(*runner_args, **kwargs):
      return self.run(None,
                      feed_dict=None,
                      options=kwargs.get("options", None),
                      run_metadata=kwargs.get("run_metadata", None),
                      callable_runner=runner,
                      callable_runner_args=runner_args)
    return wrapped_runner

  def _make_callable_from_options(self, callable_options):
    def wrapped_runner(*feed_values, **kwargs):
      return self.run(None,
                      run_metadata=kwargs.get("run_metadata", None),
                      callable_options=callable_options,
                      callable_runner_args=feed_values)
    return wrapped_runner

  @property
  def run_call_count(self):
    return self._run_call_count

  def increment_run_call_count(self):
    self._run_call_count += 1

  def _is_disk_usage_reset_each_run(self):
    """Indicates whether disk usage is reset after each Session.run.

    Subclasses that clean up the disk usage after every run should
    override this protected method.

    Returns:
      (`bool`) Whether the disk usage amount is reset to zero after
        each Session.run.
    """
    return False

  def _decorate_run_options_for_debug(
      self,
      run_options,
      debug_urls,
      debug_ops="DebugIdentity",
      node_name_regex_allowlist=None,
      op_type_regex_allowlist=None,
      tensor_dtype_regex_allowlist=None,
      tolerate_debug_op_creation_failures=False):
    """Modify a RunOptions object for debug tensor watching.

    Specifies request for outputting partition graphs. Adds
    debug_tensor_watch_opts with proper debug URLs.

    Args:
      run_options: (RunOptions) the modified RunOptions object.
      debug_urls: (list of str) debug URLs to be entered in run_options.
        debug_tensor_watch_opts.
      debug_ops: (str or list of str) debug op(s) to be used by the debugger.
      node_name_regex_allowlist: Regular-expression allowlist for node
        name.
      op_type_regex_allowlist: Regular-expression allowlist for op type.
      tensor_dtype_regex_allowlist: Regular-expression allowlist for tensor
        dtype.
      tolerate_debug_op_creation_failures: Whether debug op creation failures
        are to be tolerated.
    """

    run_options.output_partition_graphs = True
    debug_utils.watch_graph(
        run_options,
        self._sess.graph,
        debug_urls=debug_urls,
        debug_ops=debug_ops,
        node_name_regex_allowlist=node_name_regex_allowlist,
        op_type_regex_allowlist=op_type_regex_allowlist,
        tensor_dtype_regex_allowlist=tensor_dtype_regex_allowlist,
        tolerate_debug_op_creation_failures=tolerate_debug_op_creation_failures,
        reset_disk_byte_usage=(self._run_call_count == 1 or
                               self._is_disk_usage_reset_each_run()))

  def _decorate_run_options_for_profile(self, run_options):
    """Modify a RunOptions object for profiling TensorFlow graph execution.

    Args:
      run_options: (RunOptions) the modified RunOptions object.
    """

    run_options.trace_level = config_pb2.RunOptions.FULL_TRACE

  @abc.abstractmethod
  def on_session_init(self, request):
    """Callback invoked during construction of the debug-wrapper session.

    This is a blocking callback.
    The invocation happens right before the constructor ends.

    Args:
      request: (`OnSessionInitRequest`) callback request carrying information
        such as the session being wrapped.

    Returns:
      An instance of `OnSessionInitResponse`.
    """

  @abc.abstractmethod
  def on_run_start(self, request):
    """Callback invoked on run() calls to the debug-wrapper session.

    This is a blocking callback.
    The invocation happens after the wrapper's run() call is entered,
    after an increment of run call counter.

    Args:
      request: (`OnRunStartRequest`) callback request object carrying
        information about the run call such as the fetches, feed dict, run
        options, run metadata, and how many `run()` calls to this wrapper
        session have occurred.

    Returns:
      An instance of `OnRunStartResponse`, carrying information to
        debug URLs used to watch the tensors.
    """

  @abc.abstractmethod
  def on_run_end(self, request):
    """Callback invoked on run() calls to the debug-wrapper session.

    This is a blocking callback.
    The invocation happens right before the wrapper exits its run() call.

    Args:
      request: (`OnRunEndRequest`) callback request object carrying information
        such as the actual action performed by the session wrapper for the
        run() call.

    Returns:
      An instance of `OnRunStartResponse`.
    """

  def as_default(self):
    return stack.default_session(self)

  def __enter__(self):
    if self._default_session_context_manager is None:
      self._default_session_context_manager = self.as_default()
    return self._default_session_context_manager.__enter__()

  def __exit__(self, exec_type, exec_value, exec_tb):
    self._default_session_context_manager.__exit__(
        exec_type, exec_value, exec_tb)

  def __del__(self):
    if hasattr(self._sess, "__del__"):
      self._sess.__del__()

  def close(self):
    self._sess.close()

  # TODO(cais): Add _node_name_regex_allowlist and
  #   _node_op_type_regex_allowlist.

  def should_stop(self):
    if hasattr(self._sess, "should_stop"):
      return self._sess.should_stop()
    else:
      raise ValueError(
          "The wrapped session %r does not have a method called 'should_stop'. "
          "Do you intend to wrap a tf.MonitoredSession instead?" % self._sess)

class WatchOptions:
  """Type for return values of watch_fn."""

  def __init__(self,
               debug_ops=None,
               node_name_regex_allowlist=None,
               op_type_regex_allowlist=None,
               tensor_dtype_regex_allowlist=None,
               tolerate_debug_op_creation_failures=False):
    """Constructor of WatchOptions: Debug watch options.

    Used as return values of `watch_fn`s.

    Args:
      debug_ops: (`str` or `list of str`) Debug ops to be used.
      node_name_regex_allowlist: Regular-expression allowlist for node_name,
        e.g., `"(weight_[0-9]+|bias_.*)"`
      op_type_regex_allowlist: Regular-expression allowlist for the op type of
        nodes, e.g., `"(Variable|Add)"`.
        If both `node_name_regex_allowlist` and `op_type_regex_allowlist`
        are set, the two filtering operations will occur in a logical `AND`
        relation. In other words, a node will be included if and only if it
        hits both allowlists.
      tensor_dtype_regex_allowlist: Regular-expression allowlist for Tensor
        data type, e.g., `"^int.*"`.
        This allowlist operates in logical `AND` relations to the two allowlists
        above.
      tolerate_debug_op_creation_failures: (`bool`) whether debug op creation
        failures (e.g., due to dtype incompatibility) are to be tolerated by not
        throwing exceptions.
    """
    if debug_ops:
      self.debug_ops = debug_ops
    else:
      self.debug_ops = ["DebugIdentity"]
    self.node_name_regex_allowlist = node_name_regex_allowlist
    self.op_type_regex_allowlist = op_type_regex_allowlist
    self.tensor_dtype_regex_allowlist = tensor_dtype_regex_allowlist
    self.tolerate_debug_op_creation_failures = (
        tolerate_debug_op_creation_failures)

  def __repr__(self):
    return ("WatchOptions(debug_ops=%r, node_name_regex_allowlist=%r, "
            "op_type_regex_allowlist=%r, tensor_dtype_regex_allowlist=%r, "
            "tolerate_debug_op_creation_failures=%r)" %
            (self.debug_ops, self.node_name_regex_allowlist,
             self.op_type_regex_allowlist, self.tensor_dtype_regex_allowlist,
             self.tolerate_debug_op_creation_failures))

class NonInteractiveDebugWrapperSession(BaseDebugWrapperSession):
  """Base class for non-interactive (i.e., non-CLI) debug wrapper sessions."""

  def __init__(self, sess, watch_fn=None, thread_name_filter=None,
               pass_through_operrors=False):
    """Constructor of NonInteractiveDebugWrapperSession.

    Args:
      sess: The TensorFlow `Session` object being wrapped.
      watch_fn: (`Callable`) A Callable that maps the fetches and feeds of a
        debugged `Session.run()` call to `WatchOptions.`
        * Args:
          * `fetches`: the fetches to the `Session.run()` call.
          * `feeds`: the feeds to the `Session.run()` call.

        * Returns:
         (`tf_debug.WatchOptions`) An object containing debug options including
           the debug ops to use, the node names, op types and/or tensor data
           types to watch, etc. See the documentation of `tf_debug.WatchOptions`
           for more details.
      thread_name_filter: Regular-expression white list for threads on which the
        wrapper session will be active. See doc of `BaseDebugWrapperSession` for
        more details.
      pass_through_operrors: If true, all captured OpErrors will be
        propagated.  By default this captures all OpErrors.
    Raises:
       TypeError: If a non-None `watch_fn` is specified and it is not callable.
    """

    BaseDebugWrapperSession.__init__(
        self, sess, thread_name_filter=thread_name_filter,
        pass_through_operrors=pass_through_operrors)

    self._watch_fn = None
    if watch_fn is not None:
      if not callable(watch_fn):
        raise TypeError("watch_fn is not callable")
      self._watch_fn = watch_fn

  def on_session_init(self, request):
    """See doc of BaseDebugWrapperSession.on_run_start."""

    return OnSessionInitResponse(OnSessionInitAction.PROCEED)

  @abc.abstractmethod
  def prepare_run_debug_urls(self, fetches, feed_dict):
    """Abstract method to be implemented by concrete subclasses.

    This method prepares the run-specific debug URL(s).

    Args:
      fetches: Same as the `fetches` argument to `Session.run()`
      feed_dict: Same as the `feed_dict` argument to `Session.run()`

    Returns:
      debug_urls: (`str` or `list` of `str`) Debug URLs to be used in
        this `Session.run()` call.
    """

  def on_run_start(self, request):
    """See doc of BaseDebugWrapperSession.on_run_start."""

    debug_urls, watch_opts = self._prepare_run_watch_config(
        request.fetches, request.feed_dict)

    return OnRunStartResponse(
        OnRunStartAction.DEBUG_RUN,
        debug_urls,
        debug_ops=watch_opts.debug_ops,
        node_name_regex_allowlist=watch_opts.node_name_regex_allowlist,
        op_type_regex_allowlist=watch_opts.op_type_regex_allowlist,
        tensor_dtype_regex_allowlist=watch_opts.tensor_dtype_regex_allowlist,
        tolerate_debug_op_creation_failures=(
            watch_opts.tolerate_debug_op_creation_failures))

  def _prepare_run_watch_config(self, fetches, feed_dict):
    """Get the debug_urls, and node/op allowlists for the current run() call.

    Args:
      fetches: Same as the `fetches` argument to `Session.run()`.
      feed_dict: Same as the `feed_dict argument` to `Session.run()`.

    Returns:
      debug_urls: (str or list of str) Debug URLs for the current run() call.
        Currently, the list consists of only one URL that is a file:// URL.
      watch_options: (WatchOptions) The return value of a watch_fn, containing
        options including debug_ops, and allowlists.
    """

    debug_urls = self.prepare_run_debug_urls(fetches, feed_dict)
    if self._watch_fn is None:
      watch_options = WatchOptions()
    else:
      watch_options = self._watch_fn(fetches, feed_dict)
      if isinstance(watch_options, tuple):
        # For legacy return type (tuples).
        watch_options = WatchOptions(*watch_options)

    return debug_urls, watch_options

  def on_run_end(self, request):
    """See doc of BaseDebugWrapperSession.on_run_end."""

    return OnRunEndResponse()

Output: _check_type

Output: OnRunEndRequest

Output: BaseDebugWrapperSession

Output: OnSessionInitRequest

Output: WatchOptions

Output: OnSessionInitAction

Output: OnSessionInitResponse

Output: NonInteractiveDebugWrapperSession

Output: OnRunEndResponse

Output: OnRunStartRequest

Output: OnRunStartResponse

Output: OnRunStartAction


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.
# ==============================================================================
"""Tests for Building Blocks of the TensorFlow Debugger CLI."""
import os
import stat
import tempfile

import numpy as np

from tensorflow.python.client import pywrap_tf_session
from tensorflow.python.debug.cli import debugger_cli_common
from tensorflow.python.framework import test_util
from tensorflow.python.platform import gfile
from tensorflow.python.platform import googletest

class CommandLineExitTest(test_util.TensorFlowTestCase):

  def testConstructionWithoutToken(self):
    exit_exc = debugger_cli_common.CommandLineExit()

    self.assertTrue(isinstance(exit_exc, Exception))

  def testConstructionWithToken(self):
    exit_exc = debugger_cli_common.CommandLineExit(exit_token={"foo": "bar"})

    self.assertTrue(isinstance(exit_exc, Exception))
    self.assertEqual({"foo": "bar"}, exit_exc.exit_token)

class RichTextLinesTest(test_util.TensorFlowTestCase):

  def testRichTextLinesConstructorComplete(self):
    # Test RichTextLines constructor.
    screen_output = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"],
        font_attr_segs={0: [(0, 5, "red")],
                        1: [(0, 7, "blue")]},
        annotations={0: "longer wavelength",
                     1: "shorter wavelength"})

    self.assertEqual(2, len(screen_output.lines))
    self.assertEqual(2, len(screen_output.font_attr_segs))
    self.assertEqual(1, len(screen_output.font_attr_segs[0]))
    self.assertEqual(1, len(screen_output.font_attr_segs[1]))
    self.assertEqual(2, len(screen_output.annotations))

    self.assertEqual(2, screen_output.num_lines())

  def testRichTextLinesConstructorWithInvalidType(self):
    with self.assertRaisesRegex(ValueError, "Unexpected type in lines"):
      debugger_cli_common.RichTextLines(123)

  def testRichTextLinesConstructorWithString(self):
    # Test constructing a RichTextLines object with a string, instead of a list
    # of strings.
    screen_output = debugger_cli_common.RichTextLines(
        "Roses are red",
        font_attr_segs={0: [(0, 5, "red")]},
        annotations={0: "longer wavelength"})

    self.assertEqual(1, len(screen_output.lines))
    self.assertEqual(1, len(screen_output.font_attr_segs))
    self.assertEqual(1, len(screen_output.font_attr_segs[0]))
    self.assertEqual(1, len(screen_output.annotations))

  def testRichLinesAppendRichLine(self):
    rtl = debugger_cli_common.RichTextLines(
        "Roses are red",
        font_attr_segs={0: [(0, 5, "red")]})
    rtl.append_rich_line(debugger_cli_common.RichLine("Violets are ") +
                         debugger_cli_common.RichLine("blue", "blue"))
    self.assertEqual(2, len(rtl.lines))
    self.assertEqual(2, len(rtl.font_attr_segs))
    self.assertEqual(1, len(rtl.font_attr_segs[0]))
    self.assertEqual(1, len(rtl.font_attr_segs[1]))

  def testRichLineLenMethodWorks(self):
    self.assertEqual(0, len(debugger_cli_common.RichLine()))
    self.assertEqual(0, len(debugger_cli_common.RichLine("")))
    self.assertEqual(1, len(debugger_cli_common.RichLine("x")))
    self.assertEqual(6, len(debugger_cli_common.RichLine("x y z ", "blue")))

  def testRichTextLinesConstructorIncomplete(self):
    # Test RichTextLines constructor, with incomplete keyword arguments.
    screen_output = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"],
        font_attr_segs={0: [(0, 5, "red")],
                        1: [(0, 7, "blue")]})

    self.assertEqual(2, len(screen_output.lines))
    self.assertEqual(2, len(screen_output.font_attr_segs))
    self.assertEqual(1, len(screen_output.font_attr_segs[0]))
    self.assertEqual(1, len(screen_output.font_attr_segs[1]))
    self.assertEqual({}, screen_output.annotations)

  def testModifyRichTextLinesObject(self):
    screen_output = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"])

    self.assertEqual(2, len(screen_output.lines))

    screen_output.lines.append("Sugar is sweet")
    self.assertEqual(3, len(screen_output.lines))

  def testMergeRichTextLines(self):
    screen_output_1 = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"],
        font_attr_segs={0: [(0, 5, "red")],
                        1: [(0, 7, "blue")]},
        annotations={0: "longer wavelength",
                     1: "shorter wavelength"})
    screen_output_2 = debugger_cli_common.RichTextLines(
        ["Lilies are white", "Sunflowers are yellow"],
        font_attr_segs={0: [(0, 6, "white")],
                        1: [(0, 7, "yellow")]},
        annotations={
            "metadata": "foo",
            0: "full spectrum",
            1: "medium wavelength"
        })

    screen_output_1.extend(screen_output_2)

    self.assertEqual(4, screen_output_1.num_lines())
    self.assertEqual([
        "Roses are red", "Violets are blue", "Lilies are white",
        "Sunflowers are yellow"
    ], screen_output_1.lines)
    self.assertEqual({
        0: [(0, 5, "red")],
        1: [(0, 7, "blue")],
        2: [(0, 6, "white")],
        3: [(0, 7, "yellow")]
    }, screen_output_1.font_attr_segs)
    self.assertEqual({
        0: [(0, 5, "red")],
        1: [(0, 7, "blue")],
        2: [(0, 6, "white")],
        3: [(0, 7, "yellow")]
    }, screen_output_1.font_attr_segs)
    self.assertEqual({
        "metadata": "foo",
        0: "longer wavelength",
        1: "shorter wavelength",
        2: "full spectrum",
        3: "medium wavelength"
    }, screen_output_1.annotations)

  def testMergeRichTextLinesEmptyOther(self):
    screen_output_1 = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"],
        font_attr_segs={0: [(0, 5, "red")],
                        1: [(0, 7, "blue")]},
        annotations={0: "longer wavelength",
                     1: "shorter wavelength"})
    screen_output_2 = debugger_cli_common.RichTextLines([])

    screen_output_1.extend(screen_output_2)

    self.assertEqual(2, screen_output_1.num_lines())
    self.assertEqual(["Roses are red", "Violets are blue"],
                     screen_output_1.lines)
    self.assertEqual({
        0: [(0, 5, "red")],
        1: [(0, 7, "blue")],
    }, screen_output_1.font_attr_segs)
    self.assertEqual({
        0: [(0, 5, "red")],
        1: [(0, 7, "blue")],
    }, screen_output_1.font_attr_segs)
    self.assertEqual({
        0: "longer wavelength",
        1: "shorter wavelength",
    }, screen_output_1.annotations)

  def testMergeRichTextLinesEmptySelf(self):
    screen_output_1 = debugger_cli_common.RichTextLines([])
    screen_output_2 = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"],
        font_attr_segs={0: [(0, 5, "red")],
                        1: [(0, 7, "blue")]},
        annotations={0: "longer wavelength",
                     1: "shorter wavelength"})

    screen_output_1.extend(screen_output_2)

    self.assertEqual(2, screen_output_1.num_lines())
    self.assertEqual(["Roses are red", "Violets are blue"],
                     screen_output_1.lines)
    self.assertEqual({
        0: [(0, 5, "red")],
        1: [(0, 7, "blue")],
    }, screen_output_1.font_attr_segs)
    self.assertEqual({
        0: [(0, 5, "red")],
        1: [(0, 7, "blue")],
    }, screen_output_1.font_attr_segs)
    self.assertEqual({
        0: "longer wavelength",
        1: "shorter wavelength",
    }, screen_output_1.annotations)

  def testAppendALineWithAttributeSegmentsWorks(self):
    screen_output_1 = debugger_cli_common.RichTextLines(
        ["Roses are red"],
        font_attr_segs={0: [(0, 5, "red")]},
        annotations={0: "longer wavelength"})

    screen_output_1.append("Violets are blue", [(0, 7, "blue")])

    self.assertEqual(["Roses are red", "Violets are blue"],
                     screen_output_1.lines)
    self.assertEqual({
        0: [(0, 5, "red")],
        1: [(0, 7, "blue")],
    }, screen_output_1.font_attr_segs)

  def testPrependALineWithAttributeSegmentsWorks(self):
    screen_output_1 = debugger_cli_common.RichTextLines(
        ["Roses are red"],
        font_attr_segs={0: [(0, 5, "red")]},
        annotations={0: "longer wavelength"})

    screen_output_1.prepend("Violets are blue", font_attr_segs=[(0, 7, "blue")])

    self.assertEqual(["Violets are blue", "Roses are red"],
                     screen_output_1.lines)
    self.assertEqual({
        0: [(0, 7, "blue")],
        1: [(0, 5, "red")],
    }, screen_output_1.font_attr_segs)

  def testWriteToFileSucceeds(self):
    screen_output = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"],
        font_attr_segs={0: [(0, 5, "red")],
                        1: [(0, 7, "blue")]})

    fd, file_path = tempfile.mkstemp()
    os.close(fd)  # file opened exclusively, so we need to close this
    # a better fix would be to make the API take a fd
    screen_output.write_to_file(file_path)

    with gfile.Open(file_path, "r") as f:
      self.assertEqual("Roses are red\nViolets are blue\n", f.read())

    # Clean up.
    gfile.Remove(file_path)

  def testAttemptToWriteToADirectoryFails(self):
    screen_output = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"],
        font_attr_segs={0: [(0, 5, "red")],
                        1: [(0, 7, "blue")]})

    with self.assertRaises(Exception):
      screen_output.write_to_file("/")

  def testAttemptToWriteToFileInNonexistentDirectoryFails(self):
    screen_output = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"],
        font_attr_segs={0: [(0, 5, "red")],
                        1: [(0, 7, "blue")]})

    file_path = os.path.join(tempfile.mkdtemp(), "foo", "bar.txt")
    with self.assertRaises(Exception):
      screen_output.write_to_file(file_path)

class CommandHandlerRegistryTest(test_util.TensorFlowTestCase):

  def setUp(self):
    self._intentional_error_msg = "Intentionally raised exception"

  def _noop_handler(self, argv, screen_info=None):
    # A handler that does nothing other than returning "Done."
    return debugger_cli_common.RichTextLines(["Done."])

  def _handler_raising_exception(self, argv, screen_info=None):
    # A handler that intentionally raises an exception.
    raise RuntimeError(self._intentional_error_msg)

  def _handler_returning_wrong_type(self, argv, screen_info=None):
    # A handler that returns a wrong type, instead of the correct type
    # (RichTextLines).
    return "Hello"

  def _echo_screen_cols(self, argv, screen_info=None):
    # A handler that uses screen_info.
    return debugger_cli_common.RichTextLines(
        ["cols = %d" % screen_info["cols"]])

  def _exiting_handler(self, argv, screen_info=None):
    """A handler that exits with an exit token."""

    if argv:
      exit_token = argv[0]
    else:
      exit_token = None

    raise debugger_cli_common.CommandLineExit(exit_token=exit_token)

  def testRegisterEmptyCommandPrefix(self):
    registry = debugger_cli_common.CommandHandlerRegistry()

    # Attempt to register an empty-string as a command prefix should trigger
    # an exception.
    with self.assertRaisesRegex(ValueError, "Empty command prefix"):
      registry.register_command_handler("", self._noop_handler, "")

  def testRegisterAndInvokeHandler(self):
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler("noop", self._noop_handler, "")

    self.assertTrue(registry.is_registered("noop"))
    self.assertFalse(registry.is_registered("beep"))

    cmd_output = registry.dispatch_command("noop", [])
    self.assertEqual(["Done."], cmd_output.lines)

    # Attempt to invoke an unregistered command prefix should trigger an
    # exception.
    with self.assertRaisesRegex(ValueError, "No handler is registered"):
      registry.dispatch_command("beep", [])

    # Empty command prefix should trigger an exception.
    with self.assertRaisesRegex(ValueError, "Prefix is empty"):
      registry.dispatch_command("", [])

  def testExitingHandler(self):
    """Test that exit exception is correctly raised."""

    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler("exit", self._exiting_handler, "")

    self.assertTrue(registry.is_registered("exit"))

    exit_token = None
    try:
      registry.dispatch_command("exit", ["foo"])
    except debugger_cli_common.CommandLineExit as e:
      exit_token = e.exit_token

    self.assertEqual("foo", exit_token)

  def testInvokeHandlerWithScreenInfo(self):
    registry = debugger_cli_common.CommandHandlerRegistry()

    # Register and invoke a command handler that uses screen_info.
    registry.register_command_handler("cols", self._echo_screen_cols, "")

    cmd_output = registry.dispatch_command(
        "cols", [], screen_info={"cols": 100})
    self.assertEqual(["cols = 100"], cmd_output.lines)

  def testRegisterAndInvokeHandlerWithAliases(self):
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler(
        "noop", self._noop_handler, "", prefix_aliases=["n", "NOOP"])

    # is_registered() should work for full prefix and aliases.
    self.assertTrue(registry.is_registered("noop"))
    self.assertTrue(registry.is_registered("n"))
    self.assertTrue(registry.is_registered("NOOP"))

    cmd_output = registry.dispatch_command("n", [])
    self.assertEqual(["Done."], cmd_output.lines)

    cmd_output = registry.dispatch_command("NOOP", [])
    self.assertEqual(["Done."], cmd_output.lines)

  def testHandlerWithWrongReturnType(self):
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler("wrong_return",
                                      self._handler_returning_wrong_type, "")

    # If the command handler fails to return a RichTextLines instance, an error
    # should be triggered.
    with self.assertRaisesRegex(
        ValueError,
        "Return value from command handler.*is not None or a RichTextLines "
        "instance"):
      registry.dispatch_command("wrong_return", [])

  def testRegisterDuplicateHandlers(self):
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler("noop", self._noop_handler, "")

    # Registering the same command prefix more than once should trigger an
    # exception.
    with self.assertRaisesRegex(
        ValueError, "A handler is already registered for command prefix"):
      registry.register_command_handler("noop", self._noop_handler, "")

    cmd_output = registry.dispatch_command("noop", [])
    self.assertEqual(["Done."], cmd_output.lines)

  def testRegisterDuplicateAliases(self):
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler(
        "noop", self._noop_handler, "", prefix_aliases=["n"])

    # Clash with existing alias.
    with self.assertRaisesRegex(ValueError,
                                "clashes with existing prefixes or aliases"):
      registry.register_command_handler(
          "cols", self._echo_screen_cols, "", prefix_aliases=["n"])

    # The name clash should have prevent the handler from being registered.
    self.assertFalse(registry.is_registered("cols"))

    # Aliases can also clash with command prefixes.
    with self.assertRaisesRegex(ValueError,
                                "clashes with existing prefixes or aliases"):
      registry.register_command_handler(
          "cols", self._echo_screen_cols, "", prefix_aliases=["noop"])

    self.assertFalse(registry.is_registered("cols"))

  def testDispatchHandlerRaisingException(self):
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler("raise_exception",
                                      self._handler_raising_exception, "")

    # The registry should catch and wrap exceptions that occur during command
    # handling.
    cmd_output = registry.dispatch_command("raise_exception", [])
    # The error output contains a stack trace.
    # So the line count should be >= 2.
    self.assertGreater(len(cmd_output.lines), 2)
    self.assertTrue(cmd_output.lines[0].startswith(
        "Error occurred during handling of command"))
    self.assertTrue(cmd_output.lines[1].endswith(self._intentional_error_msg))

  def testRegisterNonCallableHandler(self):
    registry = debugger_cli_common.CommandHandlerRegistry()

    # Attempt to register a non-callable handler should fail.
    with self.assertRaisesRegex(ValueError, "handler is not callable"):
      registry.register_command_handler("non_callable", 1, "")

  def testRegisterHandlerWithInvalidHelpInfoType(self):
    registry = debugger_cli_common.CommandHandlerRegistry()

    with self.assertRaisesRegex(ValueError, "help_info is not a str"):
      registry.register_command_handler("noop", self._noop_handler, ["foo"])

  def testGetHelpFull(self):
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler(
        "noop",
        self._noop_handler,
        "No operation.\nI.e., do nothing.",
        prefix_aliases=["n", "NOOP"])
    registry.register_command_handler(
        "cols",
        self._echo_screen_cols,
        "Show screen width in number of columns.",
        prefix_aliases=["c"])

    help_lines = registry.get_help().lines

    # The help info should list commands in alphabetically sorted order,
    # regardless of order in which the commands are registered.
    self.assertEqual("cols", help_lines[0])
    self.assertTrue(help_lines[1].endswith("Aliases: c"))
    self.assertFalse(help_lines[2])
    self.assertTrue(help_lines[3].endswith(
        "Show screen width in number of columns."))

    self.assertFalse(help_lines[4])
    self.assertFalse(help_lines[5])

    # The default help command should appear in the help output.
    self.assertEqual("help", help_lines[6])

    self.assertEqual("noop", help_lines[12])
    self.assertTrue(help_lines[13].endswith("Aliases: n, NOOP"))
    self.assertFalse(help_lines[14])
    self.assertTrue(help_lines[15].endswith("No operation."))
    self.assertTrue(help_lines[16].endswith("I.e., do nothing."))

  def testGetHelpSingleCommand(self):
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler(
        "noop",
        self._noop_handler,
        "No operation.\nI.e., do nothing.",
        prefix_aliases=["n", "NOOP"])
    registry.register_command_handler(
        "cols",
        self._echo_screen_cols,
        "Show screen width in number of columns.",
        prefix_aliases=["c"])

    # Get help info for one of the two commands, using full prefix.
    help_lines = registry.get_help("cols").lines

    self.assertTrue(help_lines[0].endswith("cols"))
    self.assertTrue(help_lines[1].endswith("Aliases: c"))
    self.assertFalse(help_lines[2])
    self.assertTrue(help_lines[3].endswith(
        "Show screen width in number of columns."))

    # Get help info for one of the two commands, using alias.
    help_lines = registry.get_help("c").lines

    self.assertTrue(help_lines[0].endswith("cols"))
    self.assertTrue(help_lines[1].endswith("Aliases: c"))
    self.assertFalse(help_lines[2])
    self.assertTrue(help_lines[3].endswith(
        "Show screen width in number of columns."))

    # Get help info for a nonexistent command.
    help_lines = registry.get_help("foo").lines

    self.assertEqual("Invalid command prefix: \"foo\"", help_lines[0])

  def testHelpCommandWithoutIntro(self):
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler(
        "noop",
        self._noop_handler,
        "No operation.\nI.e., do nothing.",
        prefix_aliases=["n", "NOOP"])
    registry.register_command_handler(
        "cols",
        self._echo_screen_cols,
        "Show screen width in number of columns.",
        prefix_aliases=["c"])

    # Get help for all commands.
    output = registry.dispatch_command("help", [])
    self.assertEqual(["cols", "  Aliases: c", "",
                      "  Show screen width in number of columns.", "", "",
                      "help", "  Aliases: h", "", "  Print this help message.",
                      "", "", "noop", "  Aliases: n, NOOP", "",
                      "  No operation.", "  I.e., do nothing.", "", "",
                      "version", "  Aliases: ver", "",
                      "  Print the versions of TensorFlow and its key "
                      "dependencies.", "", ""],
                     output.lines)

    # Get help for one specific command prefix.
    output = registry.dispatch_command("help", ["noop"])
    self.assertEqual(["noop", "  Aliases: n, NOOP", "", "  No operation.",
                      "  I.e., do nothing."], output.lines)

    # Get help for a nonexistent command prefix.
    output = registry.dispatch_command("help", ["foo"])
    self.assertEqual(["Invalid command prefix: \"foo\""], output.lines)

  def testHelpCommandWithIntro(self):
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler(
        "noop",
        self._noop_handler,
        "No operation.\nI.e., do nothing.",
        prefix_aliases=["n", "NOOP"])

    help_intro = debugger_cli_common.RichTextLines(
        ["Introductory comments.", ""])
    registry.set_help_intro(help_intro)

    output = registry.dispatch_command("help", [])
    self.assertEqual(help_intro.lines + [
        "help", "  Aliases: h", "", "  Print this help message.", "", "",
        "noop", "  Aliases: n, NOOP", "", "  No operation.",
        "  I.e., do nothing.", "", "",
        "version", "  Aliases: ver", "",
        "  Print the versions of TensorFlow and its key dependencies.", "", ""
    ], output.lines)

class RegexFindTest(test_util.TensorFlowTestCase):

  def setUp(self):
    self._orig_screen_output = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"])

  def testRegexFindWithoutExistingFontAttrSegs(self):
    new_screen_output = debugger_cli_common.regex_find(self._orig_screen_output,
                                                       "are", "yellow")

    self.assertEqual(2, len(new_screen_output.font_attr_segs))
    self.assertEqual([(6, 9, "yellow")], new_screen_output.font_attr_segs[0])
    self.assertEqual([(8, 11, "yellow")], new_screen_output.font_attr_segs[1])

    # Check field in annotations carrying a list of matching line indices.
    self.assertEqual([0, 1], new_screen_output.annotations[
        debugger_cli_common.REGEX_MATCH_LINES_KEY])

  def testRegexFindWithExistingFontAttrSegs(self):
    # Add a font attribute segment first.
    self._orig_screen_output.font_attr_segs[0] = [(9, 12, "red")]
    self.assertEqual(1, len(self._orig_screen_output.font_attr_segs))

    new_screen_output = debugger_cli_common.regex_find(self._orig_screen_output,
                                                       "are", "yellow")
    self.assertEqual(2, len(new_screen_output.font_attr_segs))

    self.assertEqual([(6, 9, "yellow"), (9, 12, "red")],
                     new_screen_output.font_attr_segs[0])

    self.assertEqual([0, 1], new_screen_output.annotations[
        debugger_cli_common.REGEX_MATCH_LINES_KEY])

  def testRegexFindWithNoMatches(self):
    new_screen_output = debugger_cli_common.regex_find(self._orig_screen_output,
                                                       "infrared", "yellow")

    self.assertEqual({}, new_screen_output.font_attr_segs)
    self.assertEqual([], new_screen_output.annotations[
        debugger_cli_common.REGEX_MATCH_LINES_KEY])

  def testInvalidRegex(self):
    with self.assertRaisesRegex(ValueError, "Invalid regular expression"):
      debugger_cli_common.regex_find(self._orig_screen_output, "[", "yellow")

  def testRegexFindOnPrependedLinesWorks(self):
    rich_lines = debugger_cli_common.RichTextLines(["Violets are blue"])
    rich_lines.prepend(["Roses are red"])
    searched_rich_lines = debugger_cli_common.regex_find(
        rich_lines, "red", "bold")
    self.assertEqual(
        {0: [(10, 13, "bold")]}, searched_rich_lines.font_attr_segs)

    rich_lines = debugger_cli_common.RichTextLines(["Violets are blue"])
    rich_lines.prepend(["A poem"], font_attr_segs=[(0, 1, "underline")])
    searched_rich_lines = debugger_cli_common.regex_find(
        rich_lines, "poem", "italic")
    self.assertEqual(
        {0: [(0, 1, "underline"), (2, 6, "italic")]},
        searched_rich_lines.font_attr_segs)

class WrapScreenOutputTest(test_util.TensorFlowTestCase):

  def setUp(self):
    self._orig_screen_output = debugger_cli_common.RichTextLines(
        ["Folk song:", "Roses are red", "Violets are blue"],
        font_attr_segs={1: [(0, 5, "red"), (6, 9, "gray"), (10, 12, "red"),
                            (12, 13, "crimson")],
                        2: [(0, 7, "blue"), (8, 11, "gray"), (12, 14, "blue"),
                            (14, 16, "indigo")]},
        annotations={1: "longer wavelength",
                     2: "shorter wavelength"})

  def testNoActualWrapping(self):
    # Large column limit should lead to no actual wrapping.
    out, new_line_indices = debugger_cli_common.wrap_rich_text_lines(
        self._orig_screen_output, 100)

    self.assertEqual(self._orig_screen_output.lines, out.lines)
    self.assertEqual(self._orig_screen_output.font_attr_segs,
                     out.font_attr_segs)
    self.assertEqual(self._orig_screen_output.annotations, out.annotations)
    self.assertEqual(new_line_indices, [0, 1, 2])

  def testWrappingWithAttrCutoff(self):
    out, new_line_indices = debugger_cli_common.wrap_rich_text_lines(
        self._orig_screen_output, 11)

    # Add non-row-index field to out.
    out.annotations["metadata"] = "foo"

    # Check wrapped text.
    self.assertEqual(5, len(out.lines))
    self.assertEqual("Folk song:", out.lines[0])
    self.assertEqual("Roses are r", out.lines[1])
    self.assertEqual("ed", out.lines[2])
    self.assertEqual("Violets are", out.lines[3])
    self.assertEqual(" blue", out.lines[4])

    # Check wrapped font_attr_segs.
    self.assertFalse(0 in out.font_attr_segs)
    self.assertEqual([(0, 5, "red"), (6, 9, "gray"), (10, 11, "red")],
                     out.font_attr_segs[1])
    self.assertEqual([(0, 1, "red"), (1, 2, "crimson")], out.font_attr_segs[2])
    self.assertEqual([(0, 7, "blue"), (8, 11, "gray")], out.font_attr_segs[3])
    self.assertEqual([(1, 3, "blue"), (3, 5, "indigo")], out.font_attr_segs[4])

    # Check annotations.
    self.assertFalse(0 in out.annotations)
    self.assertEqual("longer wavelength", out.annotations[1])
    self.assertFalse(2 in out.annotations)
    self.assertEqual("shorter wavelength", out.annotations[3])
    self.assertFalse(4 in out.annotations)

    # Check that the non-row-index field is present in output.
    self.assertEqual("foo", out.annotations["metadata"])

    self.assertEqual(new_line_indices, [0, 1, 3])

  def testWrappingWithMultipleAttrCutoff(self):
    self._orig_screen_output = debugger_cli_common.RichTextLines(
        ["Folk song:", "Roses are red", "Violets are blue"],
        font_attr_segs={1: [(0, 12, "red")],
                        2: [(1, 16, "blue")]},
        annotations={1: "longer wavelength",
                     2: "shorter wavelength"})

    out, new_line_indices = debugger_cli_common.wrap_rich_text_lines(
        self._orig_screen_output, 5)

    # Check wrapped text.
    self.assertEqual(9, len(out.lines))
    self.assertEqual("Folk ", out.lines[0])
    self.assertEqual("song:", out.lines[1])
    self.assertEqual("Roses", out.lines[2])
    self.assertEqual(" are ", out.lines[3])
    self.assertEqual("red", out.lines[4])
    self.assertEqual("Viole", out.lines[5])
    self.assertEqual("ts ar", out.lines[6])
    self.assertEqual("e blu", out.lines[7])
    self.assertEqual("e", out.lines[8])

    # Check wrapped font_attr_segs.
    self.assertFalse(0 in out.font_attr_segs)
    self.assertFalse(1 in out.font_attr_segs)
    self.assertEqual([(0, 5, "red")], out.font_attr_segs[2])
    self.assertEqual([(0, 5, "red")], out.font_attr_segs[3])
    self.assertEqual([(0, 2, "red")], out.font_attr_segs[4])
    self.assertEqual([(1, 5, "blue")], out.font_attr_segs[5])
    self.assertEqual([(0, 5, "blue")], out.font_attr_segs[6])
    self.assertEqual([(0, 5, "blue")], out.font_attr_segs[7])
    self.assertEqual([(0, 1, "blue")], out.font_attr_segs[8])

    # Check annotations.
    self.assertFalse(0 in out.annotations)
    self.assertFalse(1 in out.annotations)
    self.assertEqual("longer wavelength", out.annotations[2])
    self.assertFalse(3 in out.annotations)
    self.assertFalse(4 in out.annotations)
    self.assertEqual("shorter wavelength", out.annotations[5])
    self.assertFalse(6 in out.annotations)
    self.assertFalse(7 in out.annotations)
    self.assertFalse(8 in out.annotations)

    self.assertEqual(new_line_indices, [0, 2, 5])

  def testWrappingInvalidArguments(self):
    with self.assertRaisesRegex(ValueError,
                                "Invalid type of input screen_output"):
      debugger_cli_common.wrap_rich_text_lines("foo", 12)

    with self.assertRaisesRegex(ValueError, "Invalid type of input cols"):
      debugger_cli_common.wrap_rich_text_lines(
          debugger_cli_common.RichTextLines(["foo", "bar"]), "12")

  def testWrappingEmptyInput(self):
    out, new_line_indices = debugger_cli_common.wrap_rich_text_lines(
        debugger_cli_common.RichTextLines([]), 10)

    self.assertEqual([], out.lines)
    self.assertEqual([], new_line_indices)

class SliceRichTextLinesTest(test_util.TensorFlowTestCase):

  def setUp(self):
    self._original = debugger_cli_common.RichTextLines(
        ["Roses are red", "Violets are blue"],
        font_attr_segs={0: [(0, 5, "red")],
                        1: [(0, 7, "blue")]},
        annotations={
            0: "longer wavelength",
            1: "shorter wavelength",
            "foo_metadata": "bar"
        })

  def testSliceBeginning(self):
    sliced = self._original.slice(0, 1)

    self.assertEqual(["Roses are red"], sliced.lines)
    self.assertEqual({0: [(0, 5, "red")]}, sliced.font_attr_segs)

    # Non-line-number metadata should be preserved.
    self.assertEqual({
        0: "longer wavelength",
        "foo_metadata": "bar"
    }, sliced.annotations)

    self.assertEqual(1, sliced.num_lines())

  def testSliceEnd(self):
    sliced = self._original.slice(1, 2)

    self.assertEqual(["Violets are blue"], sliced.lines)

    # The line index should have changed from 1 to 0.
    self.assertEqual({0: [(0, 7, "blue")]}, sliced.font_attr_segs)
    self.assertEqual({
        0: "shorter wavelength",
        "foo_metadata": "bar"
    }, sliced.annotations)

    self.assertEqual(1, sliced.num_lines())

  def testAttemptSliceWithNegativeIndex(self):
    with self.assertRaisesRegex(ValueError, "Encountered negative index"):
      self._original.slice(0, -1)

class TabCompletionRegistryTest(test_util.TensorFlowTestCase):

  def setUp(self):
    self._tc_reg = debugger_cli_common.TabCompletionRegistry()

    # Register the items in an unsorted order deliberately, to test the sorted
    # output from get_completions().
    self._tc_reg.register_tab_comp_context(
        ["print_tensor", "pt"],
        ["node_b:1", "node_b:2", "node_a:1", "node_a:2"])
    self._tc_reg.register_tab_comp_context(["node_info"],
                                           ["node_c", "node_b", "node_a"])

  def testTabCompletion(self):
    # The returned completions should have sorted order.
    self.assertEqual(
        (["node_a:1", "node_a:2", "node_b:1", "node_b:2"], "node_"),
        self._tc_reg.get_completions("print_tensor", "node_"))

    self.assertEqual((["node_a:1", "node_a:2", "node_b:1", "node_b:2"],
                      "node_"), self._tc_reg.get_completions("pt", ""))

    self.assertEqual((["node_a:1", "node_a:2"], "node_a:"),
                     self._tc_reg.get_completions("print_tensor", "node_a"))

    self.assertEqual((["node_a:1"], "node_a:1"),
                     self._tc_reg.get_completions("pt", "node_a:1"))

    self.assertEqual(([], ""),
                     self._tc_reg.get_completions("print_tensor", "node_a:3"))

    self.assertEqual((None, None), self._tc_reg.get_completions("foo", "node_"))

  def testExtendCompletionItems(self):
    self.assertEqual(
        (["node_a:1", "node_a:2", "node_b:1", "node_b:2"], "node_"),
        self._tc_reg.get_completions("print_tensor", "node_"))
    self.assertEqual((["node_a", "node_b", "node_c"], "node_"),
                     self._tc_reg.get_completions("node_info", "node_"))

    self._tc_reg.extend_comp_items("print_tensor", ["node_A:1", "node_A:2"])

    self.assertEqual((["node_A:1", "node_A:2", "node_a:1", "node_a:2",
                       "node_b:1", "node_b:2"], "node_"),
                     self._tc_reg.get_completions("print_tensor", "node_"))

    # Extending the completions for one of the context's context words should
    # have taken effect on other context words of the same context as well.
    self.assertEqual((["node_A:1", "node_A:2", "node_a:1", "node_a:2",
                       "node_b:1", "node_b:2"], "node_"),
                     self._tc_reg.get_completions("pt", "node_"))
    self.assertEqual((["node_a", "node_b", "node_c"], "node_"),
                     self._tc_reg.get_completions("node_info", "node_"))

  def testExtendCompletionItemsNonexistentContext(self):
    with self.assertRaisesRegex(KeyError,
                                "Context word \"foo\" has not been registered"):
      self._tc_reg.extend_comp_items("foo", ["node_A:1", "node_A:2"])

  def testRemoveCompletionItems(self):
    self.assertEqual(
        (["node_a:1", "node_a:2", "node_b:1", "node_b:2"], "node_"),
        self._tc_reg.get_completions("print_tensor", "node_"))
    self.assertEqual((["node_a", "node_b", "node_c"], "node_"),
                     self._tc_reg.get_completions("node_info", "node_"))

    self._tc_reg.remove_comp_items("pt", ["node_a:1", "node_a:2"])

    self.assertEqual((["node_b:1", "node_b:2"], "node_b:"),
                     self._tc_reg.get_completions("print_tensor", "node_"))
    self.assertEqual((["node_a", "node_b", "node_c"], "node_"),
                     self._tc_reg.get_completions("node_info", "node_"))

  def testRemoveCompletionItemsNonexistentContext(self):
    with self.assertRaisesRegex(KeyError,
                                "Context word \"foo\" has not been registered"):
      self._tc_reg.remove_comp_items("foo", ["node_a:1", "node_a:2"])

  def testDeregisterContext(self):
    self.assertEqual(
        (["node_a:1", "node_a:2", "node_b:1", "node_b:2"], "node_"),
        self._tc_reg.get_completions("print_tensor", "node_"))
    self.assertEqual((["node_a", "node_b", "node_c"], "node_"),
                     self._tc_reg.get_completions("node_info", "node_"))

    self._tc_reg.deregister_context(["print_tensor"])

    self.assertEqual((None, None),
                     self._tc_reg.get_completions("print_tensor", "node_"))

    # The alternative context word should be unaffected.
    self.assertEqual(
        (["node_a:1", "node_a:2", "node_b:1", "node_b:2"], "node_"),
        self._tc_reg.get_completions("pt", "node_"))

  def testDeregisterNonexistentContext(self):
    self.assertEqual(
        (["node_a:1", "node_a:2", "node_b:1", "node_b:2"], "node_"),
        self._tc_reg.get_completions("print_tensor", "node_"))
    self.assertEqual((["node_a", "node_b", "node_c"], "node_"),
                     self._tc_reg.get_completions("node_info", "node_"))

    self._tc_reg.deregister_context(["print_tensor"])

    with self.assertRaisesRegex(
        KeyError,
        "Cannot deregister unregistered context word \"print_tensor\""):
      self._tc_reg.deregister_context(["print_tensor"])

class CommandHistoryTest(test_util.TensorFlowTestCase):

  def setUp(self):
    self._fd, self._history_file_path = tempfile.mkstemp()
    self._cmd_hist = debugger_cli_common.CommandHistory(
        limit=3, history_file_path=self._history_file_path)

  def tearDown(self):
    if os.path.isfile(self._history_file_path):
      os.close(self._fd)
      os.remove(self._history_file_path)

  def _restoreFileReadWritePermissions(self, file_path):
    os.chmod(file_path,
             (stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH | stat.S_IWUSR |
              stat.S_IWGRP | stat.S_IWOTH))

  def testLookUpMostRecent(self):
    self.assertEqual([], self._cmd_hist.most_recent_n(3))

    self._cmd_hist.add_command("list_tensors")
    self._cmd_hist.add_command("node_info node_a")

    self.assertEqual(["node_info node_a"], self._cmd_hist.most_recent_n(1))
    self.assertEqual(["list_tensors", "node_info node_a"],
                     self._cmd_hist.most_recent_n(2))
    self.assertEqual(["list_tensors", "node_info node_a"],
                     self._cmd_hist.most_recent_n(3))

    self._cmd_hist.add_command("node_info node_b")

    self.assertEqual(["node_info node_b"], self._cmd_hist.most_recent_n(1))
    self.assertEqual(["node_info node_a", "node_info node_b"],
                     self._cmd_hist.most_recent_n(2))
    self.assertEqual(["list_tensors", "node_info node_a", "node_info node_b"],
                     self._cmd_hist.most_recent_n(3))
    self.assertEqual(["list_tensors", "node_info node_a", "node_info node_b"],
                     self._cmd_hist.most_recent_n(4))

    # Go over the limit.
    self._cmd_hist.add_command("node_info node_a")

    self.assertEqual(["node_info node_a"], self._cmd_hist.most_recent_n(1))
    self.assertEqual(["node_info node_b", "node_info node_a"],
                     self._cmd_hist.most_recent_n(2))
    self.assertEqual(
        ["node_info node_a", "node_info node_b", "node_info node_a"],
        self._cmd_hist.most_recent_n(3))
    self.assertEqual(
        ["node_info node_a", "node_info node_b", "node_info node_a"],
        self._cmd_hist.most_recent_n(4))

  def testLookUpPrefix(self):
    self._cmd_hist.add_command("node_info node_b")
    self._cmd_hist.add_command("list_tensors")
    self._cmd_hist.add_command("node_info node_a")

    self.assertEqual(["node_info node_b", "node_info node_a"],
                     self._cmd_hist.lookup_prefix("node_info", 10))

    self.assertEqual(["node_info node_a"], self._cmd_hist.lookup_prefix(
        "node_info", 1))

    self.assertEqual([], self._cmd_hist.lookup_prefix("print_tensor", 10))

  def testAddNonStrCommand(self):
    with self.assertRaisesRegex(
        TypeError, "Attempt to enter non-str entry to command history"):
      self._cmd_hist.add_command(["print_tensor node_a:0"])

  def testRepeatingCommandsDoNotGetLoggedRepeatedly(self):
    self._cmd_hist.add_command("help")
    self._cmd_hist.add_command("help")

    self.assertEqual(["help"], self._cmd_hist.most_recent_n(2))

  def testLoadingCommandHistoryFileObeysLimit(self):
    self._cmd_hist.add_command("help 1")
    self._cmd_hist.add_command("help 2")
    self._cmd_hist.add_command("help 3")
    self._cmd_hist.add_command("help 4")

    cmd_hist_2 = debugger_cli_common.CommandHistory(
        limit=3, history_file_path=self._history_file_path)
    self.assertEqual(["help 2", "help 3", "help 4"],
                     cmd_hist_2.most_recent_n(3))

    with open(self._history_file_path, "rt") as f:
      self.assertEqual(
          ["help 2\n", "help 3\n", "help 4\n"], f.readlines())

  def testCommandHistoryHandlesReadingIOErrorGraciously(self):
    with open(self._history_file_path, "wt") as f:
      f.write("help\n")

    # Change file to not readable by anyone.
    os.chmod(self._history_file_path, 0)

    # The creation of a CommandHistory object should not error out.
    debugger_cli_common.CommandHistory(
        limit=3, history_file_path=self._history_file_path)

    self._restoreFileReadWritePermissions(self._history_file_path)

class MenuNodeTest(test_util.TensorFlowTestCase):

  def testCommandTypeConstructorSucceeds(self):
    menu_node = debugger_cli_common.MenuItem("water flower", "water_flower")

    self.assertEqual("water flower", menu_node.caption)
    self.assertEqual("water_flower", menu_node.content)

  def testDisableWorks(self):
    menu_node = debugger_cli_common.MenuItem("water flower", "water_flower")
    self.assertTrue(menu_node.is_enabled())

    menu_node.disable()
    self.assertFalse(menu_node.is_enabled())
    menu_node.enable()
    self.assertTrue(menu_node.is_enabled())

  def testConstructAsDisabledWorks(self):
    menu_node = debugger_cli_common.MenuItem(
        "water flower", "water_flower", enabled=False)
    self.assertFalse(menu_node.is_enabled())

    menu_node.enable()
    self.assertTrue(menu_node.is_enabled())

class MenuTest(test_util.TensorFlowTestCase):

  def setUp(self):
    self.menu = debugger_cli_common.Menu()
    self.assertEqual(0, self.menu.num_items())

    self.node1 = debugger_cli_common.MenuItem("water flower", "water_flower")
    self.node2 = debugger_cli_common.MenuItem(
        "measure wavelength", "measure_wavelength")
    self.menu.append(self.node1)
    self.menu.append(self.node2)
    self.assertEqual(2, self.menu.num_items())

  def testFormatAsSingleLineWithStrItemAttrsWorks(self):
    output = self.menu.format_as_single_line(
        prefix="Menu: ", divider=", ", enabled_item_attrs="underline")
    self.assertEqual(["Menu: water flower, measure wavelength, "], output.lines)
    self.assertEqual((6, 18, [self.node1, "underline"]),
                     output.font_attr_segs[0][0])
    self.assertEqual((20, 38, [self.node2, "underline"]),
                     output.font_attr_segs[0][1])
    self.assertEqual({}, output.annotations)

  def testFormatAsSingleLineWithListItemAttrsWorks(self):
    output = self.menu.format_as_single_line(
        prefix="Menu: ", divider=", ", enabled_item_attrs=["underline", "bold"])
    self.assertEqual(["Menu: water flower, measure wavelength, "], output.lines)
    self.assertEqual((6, 18, [self.node1, "underline", "bold"]),
                     output.font_attr_segs[0][0])
    self.assertEqual((20, 38, [self.node2, "underline", "bold"]),
                     output.font_attr_segs[0][1])
    self.assertEqual({}, output.annotations)

  def testFormatAsSingleLineWithNoneItemAttrsWorks(self):
    output = self.menu.format_as_single_line(prefix="Menu: ", divider=", ")
    self.assertEqual(["Menu: water flower, measure wavelength, "], output.lines)
    self.assertEqual((6, 18, [self.node1]), output.font_attr_segs[0][0])
    self.assertEqual((20, 38, [self.node2]), output.font_attr_segs[0][1])
    self.assertEqual({}, output.annotations)

  def testInsertNode(self):
    self.assertEqual(["water flower", "measure wavelength"],
                     self.menu.captions())

    node2 = debugger_cli_common.MenuItem("write poem", "write_poem")
    self.menu.insert(1, node2)
    self.assertEqual(["water flower", "write poem", "measure wavelength"],
                     self.menu.captions())

    output = self.menu.format_as_single_line(prefix="Menu: ", divider=", ")
    self.assertEqual(["Menu: water flower, write poem, measure wavelength, "],
                     output.lines)

  def testFormatAsSingleLineWithDisabledNode(self):
    node2 = debugger_cli_common.MenuItem(
        "write poem", "write_poem", enabled=False)
    self.menu.append(node2)

    output = self.menu.format_as_single_line(
        prefix="Menu: ", divider=", ", disabled_item_attrs="bold")
    self.assertEqual(["Menu: water flower, measure wavelength, write poem, "],
                     output.lines)
    self.assertEqual((6, 18, [self.node1]), output.font_attr_segs[0][0])
    self.assertEqual((20, 38, [self.node2]), output.font_attr_segs[0][1])
    self.assertEqual((40, 50, ["bold"]), output.font_attr_segs[0][2])

class GetTensorFlowVersionLinesTest(test_util.TensorFlowTestCase):

  def testGetVersionWithoutDependencies(self):
    out = debugger_cli_common.get_tensorflow_version_lines()
    self.assertEqual(2, len(out.lines))
    self.assertEqual("TensorFlow version: %s" % pywrap_tf_session.__version__,
                     out.lines[0])

  def testGetVersionWithDependencies(self):
    out = debugger_cli_common.get_tensorflow_version_lines(True)
    self.assertIn("TensorFlow version: %s" % pywrap_tf_session.__version__,
                  out.lines)
    self.assertIn("  numpy: %s" % np.__version__, out.lines)

if __name__ == "__main__":
  googletest.main()

Output: WrapScreenOutputTest

Output: MenuNodeTest

Output: CommandLineExitTest

Output: MenuTest

Output: CommandHistoryTest

Output: CommandHandlerRegistryTest

Output: SliceRichTextLinesTest

Output: RichTextLinesTest

Output: RegexFindTest

Output: GetTensorFlowVersionLinesTest

Output: TabCompletionRegistryTest