astro py:docs

- Index
- Modules

### Navigation

- Astropy v0.3 »
- Module code »

# Source code for astropy.utils.console

    # -*- coding: utf-8 -*-
    # Licensed under a 3-clause BSD style license - see LICENSE.rst
    """
    Utilities for console input and output.
    """
    from __future__ import (absolute_import, division, print_function,
                            unicode_literals)

    import codecs
    import locale
    import re
    import math
    import multiprocessing
    import os
    import sys
    import threading
    import time

    try:
        import fcntl
        import termios
        import signal
        _CAN_RESIZE_TERMINAL = True
    except ImportError:
        _CAN_RESIZE_TERMINAL = False

    try:
        get_ipython()
    except NameError:
        OutStream = None
        IPythonIOStream = None
        stdio = sys
    else:
        try:
            from IPython.zmq.iostream import OutStream
            from IPython.utils import io
            # On Windows in particular this is necessary, as the io.stdout stream
            # in IPython gets hooked up to some pyreadline magic to handle colors
            stdio = io
            IPythonIOStream = io.IOStream
        except ImportError:
            OutStream = None
            IPythonIOStream = None
            stdio = sys

    try:
        import IPython  # pylint: disable=W0611
        # This is just to set a flag that IPython is installed at all
        _HAVE_IPYTHON = True
    except ImportError:
        _HAVE_IPYTHON = False

    from ..config import ConfigurationItem
    from ..extern import six
    from ..extern.six.moves import range

    from .misc import deprecated, isiterable

    __all__ = [
        'isatty', 'color_print', 'human_time', 'human_file_size',
        'ProgressBar', 'Spinner', 'print_code_line', 'ProgressBarOrSpinner']

    # Only use color by default on Windows if IPython is installed.
    USE_COLOR = ConfigurationItem(
        'use_color', sys.platform != 'win32' or _HAVE_IPYTHON,
        'When True, use ANSI color escape sequences when writing to the console.')
    from ..import UNICODE_OUTPUT

    IS_PY3 = sys.version_info[0] == 3

    _DEFAULT_ENCODING = 'utf-8'

    def isatty(file):
        """
        Returns `True` if `file` is a tty.

        Most built-in Python file-like objects have an `isatty` member,
        but some user-defined types may not, so this assumes those are not
        ttys.
        """
        if (multiprocessing.current_process().name != 'MainProcess' or
            threading.current_thread().getName() != 'MainThread'):
            return False

        if (OutStream is not None and
            isinstance(file, OutStream) and
            file.name == 'stdout'):
            return True
        elif hasattr(file, 'isatty'):
            return file.isatty()
        return False

    def _color_text(text, color):
        """
        Returns a string wrapped in ANSI color codes for coloring the
        text in a terminal::

            colored_text = color_text('Here is a message', 'blue')

        This won't actually effect the text until it is printed to the
        terminal.

        Parameters
        ----------
        text : str
            The string to return, bounded by the color codes.
        color : str
            An ANSI terminal color name. Must be one of:
            black, red, green, brown, blue, magenta, cyan, lightgrey,
            default, darkgrey, lightred, lightgreen, yellow, lightblue,
            lightmagenta, lightcyan, white, or '' (the empty string).
        """
        color_mapping = {
            'black': '0;30',
            'red': '0;31',
            'green': '0;32',
            'brown': '0;33',
            'blue': '0;34',
            'magenta': '0;35',
            'cyan': '0;36',
            'lightgrey': '0;37',
            'default': '0;39',
            'darkgrey': '1;30',
            'lightred': '1;31',
            'lightgreen': '1;32',
            'yellow': '1;33',
            'lightblue': '1;34',
            'lightmagenta': '1;35',
            'lightcyan': '1;36',
            'white': '1;37'}

        if sys.platform == 'win32' and OutStream is None:
            # On Windows do not colorize text unless in IPython
            return text

        color_code = color_mapping.get(color, '0;39')
        return '\033[{0}m{1}\033[0m'.format(color_code, text)

    def _decode_preferred_encoding(s):
        """Decode the supplied byte string using the preferred encoding
        for the locale (`locale.getpreferredencoding`) or, if the default encoding
        is invalid, fall back first on utf-8, then on latin-1 if the message cannot
        be decoded with utf-8.
        """

        enc = locale.getpreferredencoding()
        try:
            try:
                return s.decode(enc)
            except LookupError:
                enc = _DEFAULT_ENCODING
            return s.decode(enc)
        except UnicodeDecodeError:
            return s.decode('latin-1')

    def _write_with_fallback(s, write, fileobj):
        """Write the supplied string with the given write function like
        ``write(s)``, but use a writer for the locale's preferred encoding in case
        of a UnicodeEncodeError.  Failing that attempt to write with 'utf-8' or
        'latin-1'.
        """

        if IPythonIOStream is not None and isinstance(fileobj, IPythonIOStream):
            # If the output stream is an IPython.utils.io.IOStream object that's
            # not going to be very helpful to us since it doesn't raise any
            # exceptions when an error occurs writing to its underlying stream.
            # There's no advantage to us using IOStream.write directly though;
            # instead just write directly to its underlying stream:
            write = fileobj.stream.write

        try:
            write(s)
            return write
        except UnicodeEncodeError:
            # Let's try the next approach...
            pass

        enc = locale.getpreferredencoding()
        try:
            Writer = codecs.getwriter(enc)
        except LookupError:
            Writer = codecs.getwriter(_DEFAULT_ENCODING)

        f = Writer(fileobj)
        write = f.write

        try:
            write(s)
            return write
        except UnicodeEncodeError:
            Writer = codecs.getwriter('latin-1')
            f = Writer(fileobj)
            write = f.write

        # If this doesn't work let the exception bubble up; I'm out of ideas
        write(s)
        return write

    def color_print(*args, **kwargs):
        """
        Prints colors and styles to the terminal uses ANSI escape
        sequences.

        ::

           color_print('This is the color ', 'default', 'GREEN', 'green')

        Parameters
        ----------
        positional args : str
            The positional arguments come in pairs (*msg*, *color*), where
            *msg* is the string to display and *color* is the color to
            display it in.

            *color* is an ANSI terminal color name.  Must be one of:
            black, red, green, brown, blue, magenta, cyan, lightgrey,
            default, darkgrey, lightred, lightgreen, yellow, lightblue,
            lightmagenta, lightcyan, white, or '' (the empty string).

        file : writeable file-like object, optional
            Where to write to.  Defaults to `sys.stdout`.  If file is not
            a tty (as determined by calling its `isatty` member, if one
            exists), no coloring will be included.

        end : str, optional
            The ending of the message.  Defaults to ``\\n``.  The end will
            be printed after resetting any color or font state.
        """

        file = kwargs.get('file', stdio.stdout)

        end = kwargs.get('end', '\n')

        write = file.write
        if isatty(file) and USE_COLOR():
            for i in range(0, len(args), 2):
                msg = args[i]
                if i + 1 == len(args):
                    color = ''
                else:
                    color = args[i + 1]

                if color:
                    msg = _color_text(msg, color)

                # Some file objects support writing unicode sensibly on some Python
                # versions; if this fails try creating a writer using the locale's
                # preferred encoding. If that fails too give up.
                if not IS_PY3 and isinstance(msg, bytes):
                    msg = _decode_preferred_encoding(msg)

                write = _write_with_fallback(msg, write, file)

            write(end)
        else:
            for i in range(0, len(args), 2):
                msg = args[i]
                if not IS_PY3 and isinstance(msg, bytes):
                    # Support decoding bytes to unicode on Python 2; use the
                    # preferred encoding for the locale (which is *sometimes*
                    # sensible)
                    msg = _decode_preferred_encoding(msg)
                write(msg)
            write(end)

    def strip_ansi_codes(s):
        """
        Remove ANSI color codes from the string.
        """
        return re.sub('\033\[([0-9]+)(;[0-9]+)*m', '', s)

    def human_time(seconds):
        """
        Returns a human-friendly time string that is always exactly 6
        characters long.

        Depending on the number of seconds given, can be one of::

            1w 3d
            2d 4h
            1h 5m
            1m 4s
              15s

        Will be in color if console coloring is turned on.

        Parameters
        ----------
        seconds : int
            The number of seconds to represent

        Returns
        -------
        time : str
            A human-friendly representation of the given number of seconds
            that is always exactly 6 characters.
        """
        units = [
            ('y', 60 * 60 * 24 * 7 * 52),
            ('w', 60 * 60 * 24 * 7),
            ('d', 60 * 60 * 24),
            ('h', 60 * 60),
            ('m', 60),
            ('s', 1),
        ]

        seconds = int(seconds)

        if seconds < 60:
            return '   {0:02d}s'.format(seconds)
        for i in xrange(len(units) - 1):
            unit1, limit1 = units[i]
            unit2, limit2 = units[i + 1]
            if seconds >= limit1:
                return '{0:02d}{1}{2:02d}{3}'.format(
                    seconds // limit1, unit1,
                    (seconds % limit1) // limit2, unit2)
        return '  ~inf'

    def human_file_size(size):
        """
        Returns a human-friendly string representing a file size
        that is 2-4 characters long.

        For example, depending on the number of bytes given, can be one
        of::

            256b
            64k
            1.1G

        Parameters
        ----------
        size : int
            The size of the file (in bytes)

        Returns
        -------
        size : str
            A human-friendly representation of the size of the file
        """
        suffixes = ' kMGTPEH'
        if size == 0:
            num_scale = 0
        else:
            num_scale = int(math.floor(math.log(size) / math.log(1000)))
        if num_scale > 7:
            suffix = '?'
        else:
            suffix = suffixes[num_scale]
        num_scale = int(math.pow(1000, num_scale))
        value = size / num_scale
        str_value = str(value)
        if str_value[2] == '.':
            str_value = str_value[:2]
        else:
            str_value = str_value[:3]
        return "{0:>3s}{1}".format(str_value, suffix)

    class ProgressBar(six.Iterator):
        """
        A class to display a progress bar in the terminal.

        It is designed to be used either with the `with` statement::

            with ProgressBar(len(items)) as bar:
                for item in enumerate(items):
                    bar.update()

        or as a generator::

            for item in ProgressBar(items):
                item.process()
        """
        def __init__(self, total_or_items, file=None):
            """
            Parameters
            ----------
            total_or_items : int or sequence
                If an int, the number of increments in the process being
                tracked.  If a sequence, the items to iterate over.

            file : writable file-like object, optional
                The file to write the progress bar to.  Defaults to
                `sys.stdout`.  If `file` is not a tty (as determined by
                calling its `isatty` member, if any), the scrollbar will
                be completely silent.
            """
            if file is None:
                file = stdio.stdout

            if not isatty(file):
                self.update = self._silent_update
                self._silent = True
            else:
                self._silent = False

            if isiterable(total_or_items):
                self._items = iter(total_or_items)
                self._total = len(total_or_items)
            else:
                try:
                    self._total = int(total_or_items)
                except TypeError:
                    raise TypeError("First argument must be int or sequence")
                else:
                    self._items = iter(range(self._total))

            self._file = file
            self._start_time = time.time()

            self._should_handle_resize = (
                _CAN_RESIZE_TERMINAL and isatty(self._file))
            self._handle_resize()
            if self._should_handle_resize:
                signal.signal(signal.SIGWINCH, self._handle_resize)
                self._signal_set = True
            else:
                self._signal_set = False

            self._human_total = human_file_size(self._total)
            self.update(0)

        def _handle_resize(self, signum=None, frame=None):
            if self._should_handle_resize:
                # Don't import numpy at module level since it may not be
                # available in all contexts that this module is imported
                import numpy as np
                data = fcntl.ioctl(self._file, termios.TIOCGWINSZ, '\0' * 8)
                arr = np.fromstring(data, dtype=np.int16)
                terminal_width = arr[1]
            else:
                try:
                    terminal_width = int(os.environ.get('COLUMNS'))
                except (TypeError, ValueError):
                    terminal_width = 78
            self._bar_length = terminal_width - 37

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc_value, traceback):
            if not self._silent:
                if exc_type is None:
                    self.update(self._total)
                self._file.write('\n')
                self._file.flush()
                if self._signal_set:
                    signal.signal(signal.SIGWINCH, signal.SIG_DFL)

        def __iter__(self):
            return self

        def __next__(self):
            try:
                rv = next(self._items)
            except StopIteration:
                self.__exit__(None, None, None)
                raise
            else:
                self.update()
                return rv

        def update(self, value=None):
            """
            Update the progress bar to the given value (out of the total
            given to the constructor).
            """
            if value is None:
                value = self._current_value = self._current_value + 1
            else:
                self._current_value = value
            if self._total == 0:
                frac = 1.0
            else:
                frac = float(value) / float(self._total)

            file = self._file
            write = file.write

            bar_fill = int(float(self._bar_length) * frac)
            write('\r|')
            color_print('=' * bar_fill, 'blue', file=file, end='')
            if bar_fill < self._bar_length:
                color_print('>', 'green', file=file, end='')
                write('-' * (self._bar_length - bar_fill - 1))
            write('|')

            if value >= self._total:
                t = time.time() - self._start_time
                prefix = '     '
            elif value <= 0:
                t = None
                prefix = ''
            else:
                t = ((time.time() - self._start_time) * (1.0 - frac)) / frac
                prefix = ' ETA '
            write(' {0:>4s}/{1:>4s}'.format(
                human_file_size(value),
                self._human_total))
            write(' ({0:>6s}%)'.format('{0:.2f}'.format(frac * 100.0)))
            write(prefix)
            if t is not None:
                write(human_time(t))
            self._file.flush()

        def _silent_update(self, value=None):
            pass

        @classmethod
        def map(cls, function, items, multiprocess=False, file=None):
            """
            Does a `map` operation while displaying a progress bar with
            percentage complete.

            ::

                def work(i):
                    print(i)

                ProgressBar.map(work, range(50))

            Parameters
            ----------
            function : function
                Function to call for each step

            items : sequence
                Sequence where each element is a tuple of arguments to pass to
                *function*.

            multiprocess : bool, optional
                If `True`, use the `multiprocessing` module to distribute each
                task to a different processor core.

            file : writeable file-like object, optional
                The file to write the progress bar to.  Defaults to
                `sys.stdout`.  If `file` is not a tty (as determined by
                calling its `isatty` member, if any), the scrollbar will
                be completely silent.
            """
            results = []

            if file is None:
                file = stdio.stdout

            with cls(len(items), file=file) as bar:
                step_size = max(200, bar._bar_length)
                steps = max(int(float(len(items)) / step_size), 1)
                if not multiprocess:
                    for i, item in enumerate(items):
                        results.append(function(item))
                        if (i % steps) == 0:
                            bar.update(i)
                else:
                    p = multiprocessing.Pool()
                    for i, result in enumerate(
                        p.imap_unordered(function, items, steps)):
                        bar.update(i)
                        results.append(result)
                    p.close()
                    p.join()

            return results

        @deprecated('0.3', alternative='ProgressBar')
        @classmethod
        def iterate(cls, items, file=None):
            """
            Iterate over a sequence while indicating progress with a progress
            bar in the terminal.

            ::

                for item in ProgressBar.iterate(items):
                    pass

            Parameters
            ----------
            items : sequence
                A sequence of items to iterate over

            file : writeable file-like object, optional
                The file to write the progress bar to.  Defaults to
                `sys.stdout`.  If `file` is not a tty (as determined by
                calling its `isatty` member, if any), the scrollbar will
                be completely silent.

            Returns
            -------
            generator :
                A generator over `items`
            """
            if file is None:
                file = stdio.stdout
            return cls(items, file=file)

    class Spinner(object):
        """
        A class to display a spinner in the terminal.

        It is designed to be used with the `with` statement::

            with Spinner("Reticulating splines", "green") as s:
                for item in enumerate(items):
                    s.next()
        """
        _default_unicode_chars = "◓◑◒◐"
        _default_ascii_chars = "-/|\\"

        def __init__(self, msg, color='default', file=None, step=1,
                     chars=None):
            """
            Parameters
            ----------
            msg : str
                The message to print

            color : str, optional
                An ANSI terminal color name.  Must be one of: black, red,
                green, brown, blue, magenta, cyan, lightgrey, default,
                darkgrey, lightred, lightgreen, yellow, lightblue,
                lightmagenta, lightcyan, white.

            file : writeable file-like object, optional
                The file to write the spinner to.  Defaults to
                `sys.stdout`.  If `file` is not a tty (as determined by
                calling its `isatty` member, if any), the scrollbar will
                be completely silent.

            step : int, optional
                Only update the spinner every *step* steps

            chars : str, optional
                The character sequence to use for the spinner
            """
            if file is None:
                file = stdio.stdout

            self._msg = msg
            self._color = color
            self._file = file
            self._step = step
            if chars is None:
                if UNICODE_OUTPUT():
                    chars = self._default_unicode_chars
                else:
                    chars = self._default_ascii_chars
            self._chars = chars

            self._silent = not isatty(file)

        def _iterator(self):
            chars = self._chars
            index = 0
            file = self._file
            write = file.write
            flush = file.flush
            try_fallback = True

            while True:
                write('\r')
                color_print(self._msg, self._color, file=file, end='')
                write(' ')
                try:
                    if try_fallback:
                        write = _write_with_fallback(chars[index], write, file)
                    else:
                        write(chars[index])
                except UnicodeError:
                    # If even _write_with_fallback failed for any reason just give
                    # up on trying to use the unicode characters
                    chars = self._default_ascii_chars
                    write(chars[index])
                    try_fallback = False  # No good will come of using this again
                flush()
                yield

                for i in range(self._step):
                    yield

                index = (index + 1) % len(chars)

        def __enter__(self):
            if self._silent:
                return self._silent_iterator()
            else:
                return self._iterator()

        def __exit__(self, exc_type, exc_value, traceback):
            file = self._file
            write = file.write
            flush = file.flush

            if not self._silent:
                write('\r')
                color_print(self._msg, self._color, file=file, end='')
            if exc_type is None:
                color_print(' [Done]', 'green', file=file)
            else:
                color_print(' [Failed]', 'red', file=file)
            flush()

        def _silent_iterator(self):
            color_print(self._msg, self._color, file=self._file, end='')
            self._file.flush()

            while True:
                yield

    class ProgressBarOrSpinner(object):
        """
        A class that displays either a `ProgressBar` or `Spinner`
        depending on whether the total size of the operation is
        known or not.

        It is designed to be used with the `with` statement::

            if file.has_length():
                length = file.get_length()
            else:
                length = None
            bytes_read = 0
            with ProgressBarOrSpinner(length) as bar:
                while file.read(blocksize):
                    bytes_read += blocksize
                    bar.update(bytes_read)
        """

        def __init__(self, total, msg, color='default', file=None):
            """
            Parameters
            ----------
            total : int or None
                If an int, the number of increments in the process being
                tracked and a `ProgressBar` is displayed.  If `None`, a
                `Spinner` is displayed.

            msg : str
                The message to display above the `ProgressBar` or
                alongside the `Spinner`.

            color : str, optional
                The color of `msg`, if any.  Must be an ANSI terminal
                color name.  Must be one of: black, red, green, brown,
                blue, magenta, cyan, lightgrey, default, darkgrey,
                lightred, lightgreen, yellow, lightblue, lightmagenta,
                lightcyan, white.

            file : writable file-like object, optional
                The file to write the to.  Defaults to `sys.stdout`.  If
                `file` is not a tty (as determined by calling its `isatty`
                member, if any), only `msg` will be displayed: the
                `ProgressBar` or `Spinner` will be silent.
            """

            if file is None:
                file = stdio.stdout

            if total is None or not isatty(file):
                self._is_spinner = True
                self._obj = Spinner(msg, color=color, file=file)
            else:
                self._is_spinner = False
                color_print(msg, color, file=file)
                self._obj = ProgressBar(total, file=file)

        def __enter__(self):
            self._iter = self._obj.__enter__()
            return self

        def __exit__(self, exc_type, exc_value, traceback):
            return self._obj.__exit__(exc_type, exc_value, traceback)

        def update(self, value):
            """
            Update the progress bar to the given value (out of the total
            given to the constructor.
            """
            if self._is_spinner:
                next(self._iter)
            else:
                self._obj.update(value)

    def print_code_line(line, col=None, file=None, tabwidth=8, width=70):
        """
        Prints a line of source code, highlighting a particular character
        position in the line.  Useful for displaying the context of error
        messages.

        If the line is more than `width` characters, the line is truncated
        accordingly and '…' characters are inserted at the front and/or
        end.

        It looks like this::

            there_is_a_syntax_error_here :
                                         ^

        Parameters
        ----------
        line : unicode
            The line of code to display

        col : int, optional
            The character in the line to highlight.  `col` must be less
            than `len(line)`.

        file : writeable file-like object, optional
            Where to write to.  Defaults to `sys.stdout`.

        tabwidth : int, optional
            The number of spaces per tab (``'\\t'``) character.  Default
            is 8.  All tabs will be converted to spaces to ensure that the
            caret lines up with the correct column.

        width : int, optional
            The width of the display, beyond which the line will be
            truncated.  Defaults to 70 (this matches the default in the
            standard library's `textwrap` module).
        """

        if file is None:
            file = stdio.stdout

        write = file.write

        if col is not None:
            assert col < len(line)
            ntabs = line[:col].count('\t')
            col += ntabs * (tabwidth - 1)

        line = line.rstrip('\n')
        line = line.replace('\t', ' ' * tabwidth)

        if col is not None and col > width:
            new_col = min(width // 2, len(line) - col)
            offset = col - new_col
            line = line[offset + 1:]
            new_col = col
            col -= offset
            width = width - 3
            color_print('…', 'darkgrey', file=file, end='')

        if len(line) > width:
            write(line[:width - 1])
            color_print('…', 'darkgrey', file=file)
        else:
            write(line)
            write('\n')

        if col is not None:
            write(' ' * col)
            color_print('^', 'red', file=file)

    # The following four Getch* classes implement unbuffered character reading from
    # stdin on Windows, linux, MacOSX.  This is taken directly from ActiveState
    # Code Recipes:
    # http://code.activestate.com/recipes/134892-getch-like-unbuffered-character-reading-from-stdin/
    #

    class Getch(object):
        """Get a single character from standard input without screen echo.

        Returns
        -------
        char : str (one character)
        """
        def __init__(self):
            try:
                self.impl = _GetchWindows()
            except ImportError:
                try:
                    self.impl = _GetchMacCarbon()
                except (ImportError, AttributeError):
                    self.impl = _GetchUnix()

        def __call__(self):
            return self.impl()

    class _GetchUnix(object):
        def __init__(self):
            import tty
            import sys
            import termios  # import termios now or else you'll get the Unix
                            # version on the Mac

        def __call__(self):
            import sys
            import tty
            import termios
            fd = sys.stdin.fileno()
            old_settings = termios.tcgetattr(fd)
            try:
                tty.setraw(sys.stdin.fileno())
                ch = sys.stdin.read(1)
            finally:
                termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
            return ch

    class _GetchWindows(object):
        def __init__(self):
            import msvcrt

        def __call__(self):
            import msvcrt
            return msvcrt.getch()

    class _GetchMacCarbon(object):
        """
        A function which returns the current ASCII key that is down;
        if no ASCII key is down, the null string is returned.  The
        page http://www.mactech.com/macintosh-c/chap02-1.html was
        very helpful in figuring out how to do this.
        """
        def __init__(self):
            import Carbon
            Carbon.Evt  # see if it has this (in Unix, it doesn't)

        def __call__(self):
            import Carbon
            if Carbon.Evt.EventAvail(0x0008)[0] == 0:  # 0x0008 is the keyDownMask
                return ''
            else:
                #
                # The event contains the following info:
                # (what,msg,when,where,mod)=Carbon.Evt.GetNextEvent(0x0008)[1]
                #
                # The message (msg) contains the ASCII char which is
                # extracted with the 0x000000FF charCodeMask; this
                # number is converted to an ASCII character with chr() and
                # returned
                #
                (what, msg, when, where, mod) = Carbon.Evt.GetNextEvent(0x0008)[1]
                return chr(msg & 0x000000FF)