1# -*- coding: utf-8 -*- 2# Licensed under a 3-clause BSD style license - see LICENSE.rst 3""" 4Utilities for console input and output. 5""" 6 7import codecs 8import locale 9import re 10import math 11import multiprocessing 12import os 13import struct 14import sys 15import threading 16import time 17from concurrent.futures import ProcessPoolExecutor, as_completed 18 19try: 20 import fcntl 21 import termios 22 import signal 23 _CAN_RESIZE_TERMINAL = True 24except ImportError: 25 _CAN_RESIZE_TERMINAL = False 26 27from astropy import conf 28 29from .misc import isiterable 30from .decorators import classproperty 31 32 33__all__ = [ 34 'isatty', 'color_print', 'human_time', 'human_file_size', 35 'ProgressBar', 'Spinner', 'print_code_line', 'ProgressBarOrSpinner', 36 'terminal_size'] 37 38_DEFAULT_ENCODING = 'utf-8' 39 40 41class _IPython: 42 """Singleton class given access to IPython streams, etc.""" 43 44 @classproperty 45 def get_ipython(cls): 46 try: 47 from IPython import get_ipython 48 except ImportError: 49 pass 50 return get_ipython 51 52 @classproperty 53 def OutStream(cls): 54 if not hasattr(cls, '_OutStream'): 55 cls._OutStream = None 56 try: 57 cls.get_ipython() 58 except NameError: 59 return None 60 61 try: 62 from ipykernel.iostream import OutStream 63 except ImportError: 64 try: 65 from IPython.zmq.iostream import OutStream 66 except ImportError: 67 from IPython import version_info 68 if version_info[0] >= 4: 69 return None 70 71 try: 72 from IPython.kernel.zmq.iostream import OutStream 73 except ImportError: 74 return None 75 76 cls._OutStream = OutStream 77 78 return cls._OutStream 79 80 @classproperty 81 def ipyio(cls): 82 if not hasattr(cls, '_ipyio'): 83 try: 84 from IPython.utils import io 85 except ImportError: 86 cls._ipyio = None 87 else: 88 cls._ipyio = io 89 return cls._ipyio 90 91 @classmethod 92 def get_stream(cls, stream): 93 return getattr(cls.ipyio, stream) 94 95 96def _get_stdout(stderr=False): 97 """ 98 This utility function contains the logic to determine what streams to use 99 by default for standard out/err. 100 101 Typically this will just return `sys.stdout`, but it contains additional 102 logic for use in IPython on Windows to determine the correct stream to use 103 (usually ``IPython.util.io.stdout`` but only if sys.stdout is a TTY). 104 """ 105 106 if stderr: 107 stream = 'stderr' 108 else: 109 stream = 'stdout' 110 111 sys_stream = getattr(sys, stream) 112 return sys_stream 113 114 115def isatty(file): 116 """ 117 Returns `True` if ``file`` is a tty. 118 119 Most built-in Python file-like objects have an `isatty` member, 120 but some user-defined types may not, so this assumes those are not 121 ttys. 122 """ 123 if (multiprocessing.current_process().name != 'MainProcess' or 124 threading.current_thread().name != 'MainThread'): 125 return False 126 127 if hasattr(file, 'isatty'): 128 return file.isatty() 129 130 if _IPython.OutStream is None or (not isinstance(file, _IPython.OutStream)): 131 return False 132 133 # File is an IPython OutStream. Check whether: 134 # - File name is 'stdout'; or 135 # - File wraps a Console 136 if getattr(file, 'name', None) == 'stdout': 137 return True 138 139 if hasattr(file, 'stream'): 140 # FIXME: pyreadline has no had new release since 2015, drop it when 141 # IPython minversion is 5.x. 142 # On Windows, in IPython 2 the standard I/O streams will wrap 143 # pyreadline.Console objects if pyreadline is available; this should 144 # be considered a TTY. 145 try: 146 from pyreadline.console import Console as PyreadlineConsole 147 except ImportError: 148 return False 149 150 return isinstance(file.stream, PyreadlineConsole) 151 152 return False 153 154 155def terminal_size(file=None): 156 """ 157 Returns a tuple (height, width) containing the height and width of 158 the terminal. 159 160 This function will look for the width in height in multiple areas 161 before falling back on the width and height in astropy's 162 configuration. 163 """ 164 165 if file is None: 166 file = _get_stdout() 167 168 try: 169 s = struct.pack("HHHH", 0, 0, 0, 0) 170 x = fcntl.ioctl(file, termios.TIOCGWINSZ, s) 171 (lines, width, xpixels, ypixels) = struct.unpack("HHHH", x) 172 if lines > 12: 173 lines -= 6 174 if width > 10: 175 width -= 1 176 if lines <= 0 or width <= 0: 177 raise Exception('unable to get terminal size') 178 return (lines, width) 179 except Exception: 180 try: 181 # see if POSIX standard variables will work 182 return (int(os.environ.get('LINES')), 183 int(os.environ.get('COLUMNS'))) 184 except TypeError: 185 # fall back on configuration variables, or if not 186 # set, (25, 80) 187 lines = conf.max_lines 188 width = conf.max_width 189 if lines is None: 190 lines = 25 191 if width is None: 192 width = 80 193 return lines, width 194 195 196def _color_text(text, color): 197 """ 198 Returns a string wrapped in ANSI color codes for coloring the 199 text in a terminal:: 200 201 colored_text = color_text('Here is a message', 'blue') 202 203 This won't actually effect the text until it is printed to the 204 terminal. 205 206 Parameters 207 ---------- 208 text : str 209 The string to return, bounded by the color codes. 210 color : str 211 An ANSI terminal color name. Must be one of: 212 black, red, green, brown, blue, magenta, cyan, lightgrey, 213 default, darkgrey, lightred, lightgreen, yellow, lightblue, 214 lightmagenta, lightcyan, white, or '' (the empty string). 215 """ 216 color_mapping = { 217 'black': '0;30', 218 'red': '0;31', 219 'green': '0;32', 220 'brown': '0;33', 221 'blue': '0;34', 222 'magenta': '0;35', 223 'cyan': '0;36', 224 'lightgrey': '0;37', 225 'default': '0;39', 226 'darkgrey': '1;30', 227 'lightred': '1;31', 228 'lightgreen': '1;32', 229 'yellow': '1;33', 230 'lightblue': '1;34', 231 'lightmagenta': '1;35', 232 'lightcyan': '1;36', 233 'white': '1;37'} 234 235 if sys.platform == 'win32' and _IPython.OutStream is None: 236 # On Windows do not colorize text unless in IPython 237 return text 238 239 color_code = color_mapping.get(color, '0;39') 240 return f'\033[{color_code}m{text}\033[0m' 241 242 243def _decode_preferred_encoding(s): 244 """Decode the supplied byte string using the preferred encoding 245 for the locale (`locale.getpreferredencoding`) or, if the default encoding 246 is invalid, fall back first on utf-8, then on latin-1 if the message cannot 247 be decoded with utf-8. 248 """ 249 250 enc = locale.getpreferredencoding() 251 try: 252 try: 253 return s.decode(enc) 254 except LookupError: 255 enc = _DEFAULT_ENCODING 256 return s.decode(enc) 257 except UnicodeDecodeError: 258 return s.decode('latin-1') 259 260 261def _write_with_fallback(s, write, fileobj): 262 """Write the supplied string with the given write function like 263 ``write(s)``, but use a writer for the locale's preferred encoding in case 264 of a UnicodeEncodeError. Failing that attempt to write with 'utf-8' or 265 'latin-1'. 266 """ 267 try: 268 write(s) 269 return write 270 except UnicodeEncodeError: 271 # Let's try the next approach... 272 pass 273 274 enc = locale.getpreferredencoding() 275 try: 276 Writer = codecs.getwriter(enc) 277 except LookupError: 278 Writer = codecs.getwriter(_DEFAULT_ENCODING) 279 280 f = Writer(fileobj) 281 write = f.write 282 283 try: 284 write(s) 285 return write 286 except UnicodeEncodeError: 287 Writer = codecs.getwriter('latin-1') 288 f = Writer(fileobj) 289 write = f.write 290 291 # If this doesn't work let the exception bubble up; I'm out of ideas 292 write(s) 293 return write 294 295 296def color_print(*args, end='\n', **kwargs): 297 """ 298 Prints colors and styles to the terminal uses ANSI escape 299 sequences. 300 301 :: 302 303 color_print('This is the color ', 'default', 'GREEN', 'green') 304 305 Parameters 306 ---------- 307 positional args : str 308 The positional arguments come in pairs (*msg*, *color*), where 309 *msg* is the string to display and *color* is the color to 310 display it in. 311 312 *color* is an ANSI terminal color name. Must be one of: 313 black, red, green, brown, blue, magenta, cyan, lightgrey, 314 default, darkgrey, lightred, lightgreen, yellow, lightblue, 315 lightmagenta, lightcyan, white, or '' (the empty string). 316 317 file : writable file-like, optional 318 Where to write to. Defaults to `sys.stdout`. If file is not 319 a tty (as determined by calling its `isatty` member, if one 320 exists), no coloring will be included. 321 322 end : str, optional 323 The ending of the message. Defaults to ``\\n``. The end will 324 be printed after resetting any color or font state. 325 """ 326 327 file = kwargs.get('file', _get_stdout()) 328 329 write = file.write 330 if isatty(file) and conf.use_color: 331 for i in range(0, len(args), 2): 332 msg = args[i] 333 if i + 1 == len(args): 334 color = '' 335 else: 336 color = args[i + 1] 337 338 if color: 339 msg = _color_text(msg, color) 340 341 # Some file objects support writing unicode sensibly on some Python 342 # versions; if this fails try creating a writer using the locale's 343 # preferred encoding. If that fails too give up. 344 345 write = _write_with_fallback(msg, write, file) 346 347 write(end) 348 else: 349 for i in range(0, len(args), 2): 350 msg = args[i] 351 write(msg) 352 write(end) 353 354 355def strip_ansi_codes(s): 356 """ 357 Remove ANSI color codes from the string. 358 """ 359 return re.sub('\033\\[([0-9]+)(;[0-9]+)*m', '', s) 360 361 362def human_time(seconds): 363 """ 364 Returns a human-friendly time string that is always exactly 6 365 characters long. 366 367 Depending on the number of seconds given, can be one of:: 368 369 1w 3d 370 2d 4h 371 1h 5m 372 1m 4s 373 15s 374 375 Will be in color if console coloring is turned on. 376 377 Parameters 378 ---------- 379 seconds : int 380 The number of seconds to represent 381 382 Returns 383 ------- 384 time : str 385 A human-friendly representation of the given number of seconds 386 that is always exactly 6 characters. 387 """ 388 units = [ 389 ('y', 60 * 60 * 24 * 7 * 52), 390 ('w', 60 * 60 * 24 * 7), 391 ('d', 60 * 60 * 24), 392 ('h', 60 * 60), 393 ('m', 60), 394 ('s', 1), 395 ] 396 397 seconds = int(seconds) 398 399 if seconds < 60: 400 return f' {seconds:2d}s' 401 for i in range(len(units) - 1): 402 unit1, limit1 = units[i] 403 unit2, limit2 = units[i + 1] 404 if seconds >= limit1: 405 return '{:2d}{}{:2d}{}'.format( 406 seconds // limit1, unit1, 407 (seconds % limit1) // limit2, unit2) 408 return ' ~inf' 409 410 411def human_file_size(size): 412 """ 413 Returns a human-friendly string representing a file size 414 that is 2-4 characters long. 415 416 For example, depending on the number of bytes given, can be one 417 of:: 418 419 256b 420 64k 421 1.1G 422 423 Parameters 424 ---------- 425 size : int 426 The size of the file (in bytes) 427 428 Returns 429 ------- 430 size : str 431 A human-friendly representation of the size of the file 432 """ 433 if hasattr(size, 'unit'): 434 # Import units only if necessary because the import takes a 435 # significant time [#4649] 436 from astropy import units as u 437 size = u.Quantity(size, u.byte).value 438 439 suffixes = ' kMGTPEZY' 440 if size == 0: 441 num_scale = 0 442 else: 443 num_scale = int(math.floor(math.log(size) / math.log(1000))) 444 if num_scale > 7: 445 suffix = '?' 446 else: 447 suffix = suffixes[num_scale] 448 num_scale = int(math.pow(1000, num_scale)) 449 value = size / num_scale 450 str_value = str(value) 451 if suffix == ' ': 452 str_value = str_value[:str_value.index('.')] 453 elif str_value[2] == '.': 454 str_value = str_value[:2] 455 else: 456 str_value = str_value[:3] 457 return f"{str_value:>3s}{suffix}" 458 459 460class _mapfunc(object): 461 """ 462 A function wrapper to support ProgressBar.map(). 463 """ 464 465 def __init__(self, func): 466 self._func = func 467 468 def __call__(self, i_arg): 469 i, arg = i_arg 470 return i, self._func(arg) 471 472 473class ProgressBar: 474 """ 475 A class to display a progress bar in the terminal. 476 477 It is designed to be used either with the ``with`` statement:: 478 479 with ProgressBar(len(items)) as bar: 480 for item in enumerate(items): 481 bar.update() 482 483 or as a generator:: 484 485 for item in ProgressBar(items): 486 item.process() 487 """ 488 489 def __init__(self, total_or_items, ipython_widget=False, file=None): 490 """ 491 Parameters 492 ---------- 493 total_or_items : int or sequence 494 If an int, the number of increments in the process being 495 tracked. If a sequence, the items to iterate over. 496 497 ipython_widget : bool, optional 498 If `True`, the progress bar will display as an IPython 499 notebook widget. 500 501 file : writable file-like, optional 502 The file to write the progress bar to. Defaults to 503 `sys.stdout`. If ``file`` is not a tty (as determined by 504 calling its `isatty` member, if any, or special case hacks 505 to detect the IPython console), the progress bar will be 506 completely silent. 507 """ 508 if file is None: 509 file = _get_stdout() 510 511 if not ipython_widget and not isatty(file): 512 self.update = self._silent_update 513 self._silent = True 514 else: 515 self._silent = False 516 517 if isiterable(total_or_items): 518 self._items = iter(total_or_items) 519 self._total = len(total_or_items) 520 else: 521 try: 522 self._total = int(total_or_items) 523 except TypeError: 524 raise TypeError("First argument must be int or sequence") 525 else: 526 self._items = iter(range(self._total)) 527 528 self._file = file 529 self._start_time = time.time() 530 self._human_total = human_file_size(self._total) 531 self._ipython_widget = ipython_widget 532 533 self._signal_set = False 534 if not ipython_widget: 535 self._should_handle_resize = ( 536 _CAN_RESIZE_TERMINAL and self._file.isatty()) 537 self._handle_resize() 538 if self._should_handle_resize: 539 signal.signal(signal.SIGWINCH, self._handle_resize) 540 self._signal_set = True 541 542 self.update(0) 543 544 def _handle_resize(self, signum=None, frame=None): 545 terminal_width = terminal_size(self._file)[1] 546 self._bar_length = terminal_width - 37 547 548 def __enter__(self): 549 return self 550 551 def __exit__(self, exc_type, exc_value, traceback): 552 if not self._silent: 553 if exc_type is None: 554 self.update(self._total) 555 self._file.write('\n') 556 self._file.flush() 557 if self._signal_set: 558 signal.signal(signal.SIGWINCH, signal.SIG_DFL) 559 560 def __iter__(self): 561 return self 562 563 def __next__(self): 564 try: 565 rv = next(self._items) 566 except StopIteration: 567 self.__exit__(None, None, None) 568 raise 569 else: 570 self.update() 571 return rv 572 573 def update(self, value=None): 574 """ 575 Update progress bar via the console or notebook accordingly. 576 """ 577 578 # Update self.value 579 if value is None: 580 value = self._current_value + 1 581 self._current_value = value 582 583 # Choose the appropriate environment 584 if self._ipython_widget: 585 self._update_ipython_widget(value) 586 else: 587 self._update_console(value) 588 589 def _update_console(self, value=None): 590 """ 591 Update the progress bar to the given value (out of the total 592 given to the constructor). 593 """ 594 595 if self._total == 0: 596 frac = 1.0 597 else: 598 frac = float(value) / float(self._total) 599 600 file = self._file 601 write = file.write 602 603 if frac > 1: 604 bar_fill = int(self._bar_length) 605 else: 606 bar_fill = int(float(self._bar_length) * frac) 607 write('\r|') 608 color_print('=' * bar_fill, 'blue', file=file, end='') 609 if bar_fill < self._bar_length: 610 color_print('>', 'green', file=file, end='') 611 write('-' * (self._bar_length - bar_fill - 1)) 612 write('|') 613 614 if value >= self._total: 615 t = time.time() - self._start_time 616 prefix = ' ' 617 elif value <= 0: 618 t = None 619 prefix = '' 620 else: 621 t = ((time.time() - self._start_time) * (1.0 - frac)) / frac 622 prefix = ' ETA ' 623 write(f' {human_file_size(value):>4s}/{self._human_total:>4s}') 624 write(f' ({frac:>6.2%})') 625 write(prefix) 626 if t is not None: 627 write(human_time(t)) 628 self._file.flush() 629 630 def _update_ipython_widget(self, value=None): 631 """ 632 Update the progress bar to the given value (out of a total 633 given to the constructor). 634 635 This method is for use in the IPython notebook 2+. 636 """ 637 638 # Create and display an empty progress bar widget, 639 # if none exists. 640 if not hasattr(self, '_widget'): 641 # Import only if an IPython widget, i.e., widget in iPython NB 642 from IPython import version_info 643 if version_info[0] < 4: 644 from IPython.html import widgets 645 self._widget = widgets.FloatProgressWidget() 646 else: 647 _IPython.get_ipython() 648 from ipywidgets import widgets 649 self._widget = widgets.FloatProgress() 650 from IPython.display import display 651 652 display(self._widget) 653 self._widget.value = 0 654 655 # Calculate percent completion, and update progress bar 656 frac = (value/self._total) 657 self._widget.value = frac * 100 658 self._widget.description = f' ({frac:>6.2%})' 659 660 def _silent_update(self, value=None): 661 pass 662 663 @classmethod 664 def map(cls, function, items, multiprocess=False, file=None, step=100, 665 ipython_widget=False, multiprocessing_start_method=None): 666 """Map function over items while displaying a progress bar with percentage complete. 667 668 The map operation may run in arbitrary order on the items, but the results are 669 returned in sequential order. 670 671 :: 672 673 def work(i): 674 print(i) 675 676 ProgressBar.map(work, range(50)) 677 678 Parameters 679 ---------- 680 function : function 681 Function to call for each step 682 683 items : sequence 684 Sequence where each element is a tuple of arguments to pass to 685 *function*. 686 687 multiprocess : bool, int, optional 688 If `True`, use the `multiprocessing` module to distribute each task 689 to a different processor core. If a number greater than 1, then use 690 that number of cores. 691 692 ipython_widget : bool, optional 693 If `True`, the progress bar will display as an IPython 694 notebook widget. 695 696 file : writable file-like, optional 697 The file to write the progress bar to. Defaults to 698 `sys.stdout`. If ``file`` is not a tty (as determined by 699 calling its `isatty` member, if any), the scrollbar will 700 be completely silent. 701 702 step : int, optional 703 Update the progress bar at least every *step* steps (default: 100). 704 If ``multiprocess`` is `True`, this will affect the size 705 of the chunks of ``items`` that are submitted as separate tasks 706 to the process pool. A large step size may make the job 707 complete faster if ``items`` is very long. 708 709 multiprocessing_start_method : str, optional 710 Useful primarily for testing; if in doubt leave it as the default. 711 When using multiprocessing, certain anomalies occur when starting 712 processes with the "spawn" method (the only option on Windows); 713 other anomalies occur with the "fork" method (the default on 714 Linux). 715 """ 716 717 if multiprocess: 718 function = _mapfunc(function) 719 items = list(enumerate(items)) 720 721 results = cls.map_unordered( 722 function, items, multiprocess=multiprocess, 723 file=file, step=step, 724 ipython_widget=ipython_widget, 725 multiprocessing_start_method=multiprocessing_start_method) 726 727 if multiprocess: 728 _, results = zip(*sorted(results)) 729 results = list(results) 730 731 return results 732 733 @classmethod 734 def map_unordered(cls, function, items, multiprocess=False, file=None, 735 step=100, ipython_widget=False, 736 multiprocessing_start_method=None): 737 """Map function over items, reporting the progress. 738 739 Does a `map` operation while displaying a progress bar with 740 percentage complete. The map operation may run on arbitrary order 741 on the items, and the results may be returned in arbitrary order. 742 743 :: 744 745 def work(i): 746 print(i) 747 748 ProgressBar.map(work, range(50)) 749 750 Parameters 751 ---------- 752 function : function 753 Function to call for each step 754 755 items : sequence 756 Sequence where each element is a tuple of arguments to pass to 757 *function*. 758 759 multiprocess : bool, int, optional 760 If `True`, use the `multiprocessing` module to distribute each task 761 to a different processor core. If a number greater than 1, then use 762 that number of cores. 763 764 ipython_widget : bool, optional 765 If `True`, the progress bar will display as an IPython 766 notebook widget. 767 768 file : writable file-like, optional 769 The file to write the progress bar to. Defaults to 770 `sys.stdout`. If ``file`` is not a tty (as determined by 771 calling its `isatty` member, if any), the scrollbar will 772 be completely silent. 773 774 step : int, optional 775 Update the progress bar at least every *step* steps (default: 100). 776 If ``multiprocess`` is `True`, this will affect the size 777 of the chunks of ``items`` that are submitted as separate tasks 778 to the process pool. A large step size may make the job 779 complete faster if ``items`` is very long. 780 781 multiprocessing_start_method : str, optional 782 Useful primarily for testing; if in doubt leave it as the default. 783 When using multiprocessing, certain anomalies occur when starting 784 processes with the "spawn" method (the only option on Windows); 785 other anomalies occur with the "fork" method (the default on 786 Linux). 787 """ 788 789 results = [] 790 791 if file is None: 792 file = _get_stdout() 793 794 with cls(len(items), ipython_widget=ipython_widget, file=file) as bar: 795 if bar._ipython_widget: 796 chunksize = step 797 else: 798 default_step = max(int(float(len(items)) / bar._bar_length), 1) 799 chunksize = min(default_step, step) 800 if not multiprocess or multiprocess < 1: 801 for i, item in enumerate(items): 802 results.append(function(item)) 803 if (i % chunksize) == 0: 804 bar.update(i) 805 else: 806 ctx = multiprocessing.get_context(multiprocessing_start_method) 807 kwargs = dict(mp_context=ctx) 808 809 with ProcessPoolExecutor( 810 max_workers=(int(multiprocess) 811 if multiprocess is not True 812 else None), 813 **kwargs) as p: 814 for i, f in enumerate( 815 as_completed( 816 p.submit(function, item) 817 for item in items)): 818 bar.update(i) 819 results.append(f.result()) 820 821 return results 822 823 824class Spinner: 825 """ 826 A class to display a spinner in the terminal. 827 828 It is designed to be used with the ``with`` statement:: 829 830 with Spinner("Reticulating splines", "green") as s: 831 for item in enumerate(items): 832 s.update() 833 """ 834 _default_unicode_chars = "◓◑◒◐" 835 _default_ascii_chars = "-/|\\" 836 837 def __init__(self, msg, color='default', file=None, step=1, 838 chars=None): 839 """ 840 Parameters 841 ---------- 842 msg : str 843 The message to print 844 845 color : str, optional 846 An ANSI terminal color name. Must be one of: black, red, 847 green, brown, blue, magenta, cyan, lightgrey, default, 848 darkgrey, lightred, lightgreen, yellow, lightblue, 849 lightmagenta, lightcyan, white. 850 851 file : writable file-like, optional 852 The file to write the spinner to. Defaults to 853 `sys.stdout`. If ``file`` is not a tty (as determined by 854 calling its `isatty` member, if any, or special case hacks 855 to detect the IPython console), the spinner will be 856 completely silent. 857 858 step : int, optional 859 Only update the spinner every *step* steps 860 861 chars : str, optional 862 The character sequence to use for the spinner 863 """ 864 865 if file is None: 866 file = _get_stdout() 867 868 self._msg = msg 869 self._color = color 870 self._file = file 871 self._step = step 872 if chars is None: 873 if conf.unicode_output: 874 chars = self._default_unicode_chars 875 else: 876 chars = self._default_ascii_chars 877 self._chars = chars 878 879 self._silent = not isatty(file) 880 881 if self._silent: 882 self._iter = self._silent_iterator() 883 else: 884 self._iter = self._iterator() 885 886 def _iterator(self): 887 chars = self._chars 888 index = 0 889 file = self._file 890 write = file.write 891 flush = file.flush 892 try_fallback = True 893 894 while True: 895 write('\r') 896 color_print(self._msg, self._color, file=file, end='') 897 write(' ') 898 try: 899 if try_fallback: 900 write = _write_with_fallback(chars[index], write, file) 901 else: 902 write(chars[index]) 903 except UnicodeError: 904 # If even _write_with_fallback failed for any reason just give 905 # up on trying to use the unicode characters 906 chars = self._default_ascii_chars 907 write(chars[index]) 908 try_fallback = False # No good will come of using this again 909 flush() 910 yield 911 912 for i in range(self._step): 913 yield 914 915 index = (index + 1) % len(chars) 916 917 def __enter__(self): 918 return self 919 920 def __exit__(self, exc_type, exc_value, traceback): 921 file = self._file 922 write = file.write 923 flush = file.flush 924 925 if not self._silent: 926 write('\r') 927 color_print(self._msg, self._color, file=file, end='') 928 if exc_type is None: 929 color_print(' [Done]', 'green', file=file) 930 else: 931 color_print(' [Failed]', 'red', file=file) 932 flush() 933 934 def __iter__(self): 935 return self 936 937 def __next__(self): 938 next(self._iter) 939 940 def update(self, value=None): 941 """Update the spin wheel in the terminal. 942 943 Parameters 944 ---------- 945 value : int, optional 946 Ignored (present just for compatibility with `ProgressBar.update`). 947 948 """ 949 950 next(self) 951 952 def _silent_iterator(self): 953 color_print(self._msg, self._color, file=self._file, end='') 954 self._file.flush() 955 956 while True: 957 yield 958 959 960class ProgressBarOrSpinner: 961 """ 962 A class that displays either a `ProgressBar` or `Spinner` 963 depending on whether the total size of the operation is 964 known or not. 965 966 It is designed to be used with the ``with`` statement:: 967 968 if file.has_length(): 969 length = file.get_length() 970 else: 971 length = None 972 bytes_read = 0 973 with ProgressBarOrSpinner(length) as bar: 974 while file.read(blocksize): 975 bytes_read += blocksize 976 bar.update(bytes_read) 977 """ 978 979 def __init__(self, total, msg, color='default', file=None): 980 """ 981 Parameters 982 ---------- 983 total : int or None 984 If an int, the number of increments in the process being 985 tracked and a `ProgressBar` is displayed. If `None`, a 986 `Spinner` is displayed. 987 988 msg : str 989 The message to display above the `ProgressBar` or 990 alongside the `Spinner`. 991 992 color : str, optional 993 The color of ``msg``, if any. Must be an ANSI terminal 994 color name. Must be one of: black, red, green, brown, 995 blue, magenta, cyan, lightgrey, default, darkgrey, 996 lightred, lightgreen, yellow, lightblue, lightmagenta, 997 lightcyan, white. 998 999 file : writable file-like, optional 1000 The file to write the to. Defaults to `sys.stdout`. If 1001 ``file`` is not a tty (as determined by calling its `isatty` 1002 member, if any), only ``msg`` will be displayed: the 1003 `ProgressBar` or `Spinner` will be silent. 1004 """ 1005 1006 if file is None: 1007 file = _get_stdout() 1008 1009 if total is None or not isatty(file): 1010 self._is_spinner = True 1011 self._obj = Spinner(msg, color=color, file=file) 1012 else: 1013 self._is_spinner = False 1014 color_print(msg, color, file=file) 1015 self._obj = ProgressBar(total, file=file) 1016 1017 def __enter__(self): 1018 return self 1019 1020 def __exit__(self, exc_type, exc_value, traceback): 1021 return self._obj.__exit__(exc_type, exc_value, traceback) 1022 1023 def update(self, value): 1024 """ 1025 Update the progress bar to the given value (out of the total 1026 given to the constructor. 1027 """ 1028 self._obj.update(value) 1029 1030 1031def print_code_line(line, col=None, file=None, tabwidth=8, width=70): 1032 """ 1033 Prints a line of source code, highlighting a particular character 1034 position in the line. Useful for displaying the context of error 1035 messages. 1036 1037 If the line is more than ``width`` characters, the line is truncated 1038 accordingly and '…' characters are inserted at the front and/or 1039 end. 1040 1041 It looks like this:: 1042 1043 there_is_a_syntax_error_here : 1044 ^ 1045 1046 Parameters 1047 ---------- 1048 line : unicode 1049 The line of code to display 1050 1051 col : int, optional 1052 The character in the line to highlight. ``col`` must be less 1053 than ``len(line)``. 1054 1055 file : writable file-like, optional 1056 Where to write to. Defaults to `sys.stdout`. 1057 1058 tabwidth : int, optional 1059 The number of spaces per tab (``'\\t'``) character. Default 1060 is 8. All tabs will be converted to spaces to ensure that the 1061 caret lines up with the correct column. 1062 1063 width : int, optional 1064 The width of the display, beyond which the line will be 1065 truncated. Defaults to 70 (this matches the default in the 1066 standard library's `textwrap` module). 1067 """ 1068 1069 if file is None: 1070 file = _get_stdout() 1071 1072 if conf.unicode_output: 1073 ellipsis = '…' 1074 else: 1075 ellipsis = '...' 1076 1077 write = file.write 1078 1079 if col is not None: 1080 if col >= len(line): 1081 raise ValueError('col must be less the the line length.') 1082 ntabs = line[:col].count('\t') 1083 col += ntabs * (tabwidth - 1) 1084 1085 line = line.rstrip('\n') 1086 line = line.replace('\t', ' ' * tabwidth) 1087 1088 if col is not None and col > width: 1089 new_col = min(width // 2, len(line) - col) 1090 offset = col - new_col 1091 line = line[offset + len(ellipsis):] 1092 width -= len(ellipsis) 1093 new_col = col 1094 col -= offset 1095 color_print(ellipsis, 'darkgrey', file=file, end='') 1096 1097 if len(line) > width: 1098 write(line[:width - len(ellipsis)]) 1099 color_print(ellipsis, 'darkgrey', file=file) 1100 else: 1101 write(line) 1102 write('\n') 1103 1104 if col is not None: 1105 write(' ' * col) 1106 color_print('^', 'red', file=file) 1107 1108 1109# The following four Getch* classes implement unbuffered character reading from 1110# stdin on Windows, linux, MacOSX. This is taken directly from ActiveState 1111# Code Recipes: 1112# http://code.activestate.com/recipes/134892-getch-like-unbuffered-character-reading-from-stdin/ 1113# 1114 1115class Getch: 1116 """Get a single character from standard input without screen echo. 1117 1118 Returns 1119 ------- 1120 char : str (one character) 1121 """ 1122 1123 def __init__(self): 1124 try: 1125 self.impl = _GetchWindows() 1126 except ImportError: 1127 try: 1128 self.impl = _GetchMacCarbon() 1129 except (ImportError, AttributeError): 1130 self.impl = _GetchUnix() 1131 1132 def __call__(self): 1133 return self.impl() 1134 1135 1136class _GetchUnix: 1137 def __init__(self): 1138 import tty # pylint: disable=W0611 1139 import sys # pylint: disable=W0611 1140 1141 # import termios now or else you'll get the Unix 1142 # version on the Mac 1143 import termios # pylint: disable=W0611 1144 1145 def __call__(self): 1146 import sys 1147 import tty 1148 import termios 1149 fd = sys.stdin.fileno() 1150 old_settings = termios.tcgetattr(fd) 1151 try: 1152 tty.setraw(sys.stdin.fileno()) 1153 ch = sys.stdin.read(1) 1154 finally: 1155 termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) 1156 return ch 1157 1158 1159class _GetchWindows: 1160 def __init__(self): 1161 import msvcrt # pylint: disable=W0611 1162 1163 def __call__(self): 1164 import msvcrt 1165 return msvcrt.getch() 1166 1167 1168class _GetchMacCarbon: 1169 """ 1170 A function which returns the current ASCII key that is down; 1171 if no ASCII key is down, the null string is returned. The 1172 page http://www.mactech.com/macintosh-c/chap02-1.html was 1173 very helpful in figuring out how to do this. 1174 """ 1175 1176 def __init__(self): 1177 import Carbon 1178 Carbon.Evt # see if it has this (in Unix, it doesn't) 1179 1180 def __call__(self): 1181 import Carbon 1182 if Carbon.Evt.EventAvail(0x0008)[0] == 0: # 0x0008 is the keyDownMask 1183 return '' 1184 else: 1185 # 1186 # The event contains the following info: 1187 # (what,msg,when,where,mod)=Carbon.Evt.GetNextEvent(0x0008)[1] 1188 # 1189 # The message (msg) contains the ASCII char which is 1190 # extracted with the 0x000000FF charCodeMask; this 1191 # number is converted to an ASCII character with chr() and 1192 # returned 1193 # 1194 (what, msg, when, where, mod) = Carbon.Evt.GetNextEvent(0x0008)[1] 1195 return chr(msg & 0x000000FF) 1196