python-startup/python_startup.py

# Copyright (C) 2015  Stefano Mazzucco <stefano - at - curso - dot - re>
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import sys

import rlcompleter
import readline
readline.parse_and_bind('tab: complete')
del rlcompleter, readline

# ANSI colors
END_COLOR = '\033[0m'
YELLOW = '\033[33m'
RED = '\033[31m'

try:
    import builtins as __builtins__
except ImportError:
    # Python 2
    import __builtin__ as __builtins__


class PromptCounter(object):

    def __init__(self, in_fmt, out_fmt):
        self.count = 1
        self._in_fmt = in_fmt
        self._out_fmt = out_fmt

    def __str__(self):
        current = self._in_fmt.format(self.count)
        self.count += 1
        return current

    def __call__(self, value):
        # Only useful in sys.displayhook
        # https://docs.python.org/3/library/sys.html#sys.displayhook
        if value is None:
            return
        # Set '_' to None to avoid recursion
        __builtins__._ = None
        # __call__ is called before __str__ so, format with count
        # decremented by 1 to keep the two in sync
        sys.stdout.write(''.join(
            [self._out_fmt.format(self.count - 1), repr(value), '\n']))
        __builtins__._ = value

if __name__ == '__main__':
    if not os.environ.get('INSIDE_EMACS'):
        sys.ps1 = sys.displayhook = PromptCounter(
            YELLOW + '[In  {}]:' + END_COLOR + ' ',
            RED + '[Out {}]:' + END_COLOR + ' ')

    del YELLOW, RED, END_COLOR, PromptCounter