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