Smart prompts for iPython / Loggi

ipython_smart_prompt.py

# encoding: utf-8
from __future__ import absolute_import, print_function, unicode_literals

import os
from collections import defaultdict

from IPython import get_ipython
from IPython.terminal.prompts import Prompts
from django.conf import settings
from pygments.token import Token

UNEXPECTED_ALERT = ' (!)'
ENV_TOKEN_CONFIG = defaultdict(lambda: Token.Tilde, dict(
    production=Token.OutPromptNum,
    staging=Token.Digraph,
    sandbox=Token.Digraph,
    dev=Token.Prompt,
))
OUT_TOKEN_CONFIG = defaultdict(lambda: Token.OutPrompt,
                               {Token.PromptNum: Token.OutPromptNum})


def detect_django_settings():
    return settings.CURRENT_ENV_NAME


def detect_expected_settings():
    return os.getenv('APP_ENV', 'dev')


def detect_hostname():
    return os.getenv('HOSTNAME')


def detect_config_anomaly():
    return detect_expected_settings() != detect_django_settings()


class LoggiEnvPrompts(Prompts):
    def in_prompt_tokens(self, cli=None):
        anomaly = (Token.Tilde, UNEXPECTED_ALERT) if detect_config_anomaly() else None
        env_token = ENV_TOKEN_CONFIG.get(detect_django_settings(), Token.Tilde)
        return filter(None, [
            (env_token, detect_django_settings()),
            anomaly,
            (Token.Prompt, ' ['),
            (Token.PromptNum, str(self.shell.execution_count)),
            (Token.Prompt, ']: ')])

    def out_prompt_tokens(self):
        out = [(OUT_TOKEN_CONFIG[token], text)
               for token, text in self.in_prompt_tokens()]
        return out


def install():
    import time
    while True:
        ipython = get_ipython()
        if ipython:
            break
        time.sleep(.1)
    ipython.prompts = LoggiEnvPrompts(ipython)


import threading

threading.Thread(target=install).start()