pythonrc

from types import BuiltinFunctionType, FunctionType, MethodType, ModuleType
from code import InteractiveConsole
from platform import python_version
from math import sqrt, sin, cos, tan, factorial, gcd, log, e, pi
from random import randint
import subprocess
import rlcompleter
import readline
import sys

# TODO:
#   - Add support for command history

def factors(n):
    return [(i, n // i) for i in range(1, n+1) if n // i == n / i]

def split_by_sep(s, sep):
    index = s.find(sep)
    if index == -1:
        return (s, "")
    return (s[:index], s[index:].strip())

def simplify_float(f):
    if f // 1 == f:
        return int(f)
    return f

RED = "\x01\033[31m\x02"
GREEN = "\x01\033[32m\x02"
BLUE = "\x01\033[34m\x02"
PURPLE = "\x01\033[35m\x02"
CYAN = "\x01\033[36m\x02"
RESET = "\x01\033[0m\x02"

INDENT = " " * 2

def color_format(color, s):
    return f"{color}{s}{RESET}"

def pretty_display(console, value, _indent=""):
    if console.file_line != -1:
        return ""
    console.locals['_'] = value
    expand = len(str(value)) > 80
    spacing = "\n" if expand else " "
    indent = (INDENT if expand else "") + (_indent if expand else "")
    if value is None:
        return color_format(PURPLE, "None")
    elif isinstance(value, bool):
        return color_format(PURPLE, value)
    elif isinstance(value, int):
        return color_format(GREEN, value)
    elif isinstance(value, float):
        return color_format(GREEN, value)
    elif isinstance(value, complex):
        real = simplify_float(value.real)
        imag = simplify_float(value.imag)
        if real == 0:
            return color_format(BLUE, str(imag) + "i")
        elif imag == 0:
            return color_format(GREEN, str(real))
        else:
            return (color_format(CYAN, "(") +
                    color_format(GREEN, real) +
                    " " +
                    color_format(PURPLE, "+") +
                    " " +
                    color_format(BLUE, str(imag) + "i") +
                    color_format(CYAN, ")"))
    elif isinstance(value, str):
        return color_format(GREEN, repr(value))
    elif isinstance(value, list) or isinstance(value, tuple):
        obracket = "[" if isinstance(value, list) else "("
        cbracket = "]" if isinstance(value, list) else ")"
        if len(value) == 0:
            return obracket + cbracket
        output = obracket
        for i, c in enumerate(value):
            output += spacing
            output += indent
            output += pretty_display(console, c, indent)
            if i != len(value) - 1:
                output += ","
        output += spacing
        output += (_indent if expand else "")
        output += cbracket
        return output
    elif isinstance(value, dict):
        if len(value) == 0:
            return "{}"
        output = "{"
        for i, c in enumerate(value):
            output += spacing
            output += indent
            output += pretty_display(console, c, indent)
            output += " : "
            output += pretty_display(console, value[c], indent)
            if i != len(value) - 1:
                output += ","
        output += spacing
        output += (_indent if expand else "")
        output += "}"
        return output
    elif isinstance(value, type):
        return color_format(PURPLE, "type ") + color_format(GREEN, value.__name__)
    elif isinstance(value, FunctionType):
        return color_format(PURPLE, "function ") + color_format(GREEN, value.__name__)
    elif isinstance(value, BuiltinFunctionType):
        return color_format(PURPLE, "built-in function ") + color_format(GREEN, value.__name__)
    elif isinstance(value, MethodType):
        return color_format(PURPLE, "method ") + color_format(GREEN, value.__name__)
    elif isinstance(value, ModuleType):
        return color_format(PURPLE, "module ") + color_format(GREEN, value.__name__)
    elif type(value).__repr__ != object.__repr__:
        return repr(value)
    elif type(value).__str__ != object.__str__:
        return str(value)
    elif hasattr(value, "__name__"):
        return (color_format(PURPLE, "var ") +
                color_format(GREEN, value.__name__) +
                color_format(PURPLE, " of type ") +
                color_format(GREEN, type(value).__name__))
    return (color_format(PURPLE, "object of type ") +
            color_format(GREEN, type(value).__name__))

class CustomConsole(InteractiveConsole, rlcompleter.Completer):
    def __init__(self, commands):
        super(CustomConsole, self).__init__()
        self.commands = commands
        self.pretty_print = True
        self.file_line = -1
        self.file_name = ""
        self.error_flag = False
        self.init_prompt()
        self.init_display()
        self.init_readline()

    def init_prompt(self):
        sys.ps1 = color_format(BLUE, "> ")
        sys.ps2 = color_format(GREEN, ". ")

    def init_display(self):
        if not self.pretty_print:
            sys.displayhook = print
            return
        def displayhook(value):
            prettied = pretty_display(self, value)
            if prettied:
                print(prettied)
        sys.displayhook = displayhook

    def init_readline(self):
        readline.parse_and_bind("tab: complete")
        readline.read_init_file()

    def showtraceback(self, *args):
        self.write(RED)
        result = super(CustomConsole, self).showtraceback(*args)
        self.write(RESET)
        return result

    def showsyntaxerror(self, filename=None):
        self.write(RED)
        result = super(CustomConsole, self).showsyntaxerror(filename)
        self.write(RESET)
        return result

    def error(self, message):
        self.error_flag = True
        prefix = f"Error at line {self.file_line} when loading file '{self.file_name}': " if self.file_line != -1 else ""
        self.write(color_format(RED, prefix + message) + "\n")

    def raw_input(self, prompt=''):
        line = super(CustomConsole, self).raw_input(prompt)
        return self.handle(line)

    def handle(self, line):
        if len(line) and line[0] == "%":
            command, inp = split_by_sep(line, " ")
            command = command[1:]
            if command in self.commands:
                func = self.commands[command]
                func(self, inp)
                return ''
            self.error(f"Command '{command}' does not exist")
            return ''
        elif len(line.strip()) and line.strip()[-1] == "?":
            inp = line.strip()[:-1]
            try:
                value = eval(inp, {}, self.locals)
            except:
                self.error(f"Could not display docstring of value '{inp}'\n")
                return ''
            if value.__doc__ is not None:
                self.write(value.__doc__ + "\n")
            else:
                help(value)
                # self.write(color_format(RED, "Value does not have a doc string\n"))
            return ''
        elif len(line) and line[0] == "!":
            command = line[1:]
            command = command.format(**self.locals)
            status, output = subprocess.getstatusoutput(command)
            if status:
                self.error(output)
            else:
                self.write(output + "\n")
            return ''
        return line

    def interact(self):
        banner = f"Interactive repl using Python {python_version()}\nType '%help' for command usage."
        super(CustomConsole, self).interact(banner=banner, exitmsg="")
        sys.exit()

def show_help(console, _):
    """Displays this help page."""
    output = color_format(BLUE, "Commands and usage:\n")
    for command in console.commands:
        output += (color_format(PURPLE, command) +
            " : " +
            color_format(GREEN, console.commands[command].__doc__) +
            "\n")
    console.write(output)

def ptoggle(console, _):
    """Toggles pretty print functionality."""
    console.pretty_print = not console.pretty_print
    console.init_display()

def inspect(console, inp):
    """Print information about a value."""
    try:
        value = eval(inp, {}, console.locals)
    except:
        console.error(f"Could not inspect value '{inp}'")
        return
    pretty_print_types = [bool, int, str, list, tuple, dict, type, FunctionType, BuiltinFunctionType, MethodType, ModuleType]
    function_types = [FunctionType, BuiltinFunctionType, MethodType]
    if value is None or type(value) in pretty_print_types:
        console.write(pretty_display(console, value) + "\n")
        return
    _type = type(value).__name__
    attributes = [a for a in dir(value) if a[0] != "_" and type(getattr(value, a)) not in function_types]
    attributes.reverse()
    if len(attributes) == 0:
        console.write(color_format(GREEN, _type) + " {}\n")
        return
    output = color_format(GREEN, _type) + " {\n"
    for i, a in enumerate(attributes):
        output += INDENT
        output += color_format(BLUE, a)
        output += ": "
        output += pretty_display(console, getattr(value, a), INDENT)
        if i != len(attributes) - 1:
            output += ","
        output += "\n"
    output += "}\n"
    console.write(output)

def methods(console, inp):
    """Displays all methods in an object."""
    try:
        value = eval(inp, {}, console.locals)
    except:
        console.error(f"Invalid input '{inp}'")
        return
    function_types = [FunctionType, BuiltinFunctionType, MethodType]
    _type = type(value).__name__
    attributes = [a for a in dir(value) if a[0] != "_" and type(getattr(value, a)) in function_types]
    attributes.reverse()
    if len(attributes) == 0:
        console.write(color_format(GREEN, _type) + " {}\n")
        return
    output = color_format(GREEN, _type) + " {\n"
    for i, a in enumerate(attributes):
        output += INDENT
        output += color_format(BLUE, a)
        output += ": "
        output += pretty_display(console, getattr(value, a), INDENT)
        if i != len(attributes) - 1:
            output += ","
        output += "\n"
    output += "}\n"
    console.write(output)

def create_type(console, inp):
    """Creates a simple struct type using named arguments (syntax: %type TypeName param1 param2 param3 ...)"""
    words = inp.strip().split(' ')
    if len(words) <= 1:
        console.error("Not enough arguments")
        return
    type_name, *parameters = words
    if not type_name.isidentifier():
        console.error("Invalid type name")
        return
    for p in parameters:
        if not p.isidentifier():
            console.error(f"Invalid parameter name '{p}'")
            return
    class InnerType:
        def __init__(self, *args):
            if len(args) != len(parameters):
                raise TypeError(f"Invalid amount of arguments passed to struct '{type_name}'")
            for i, name in enumerate(parameters):
                setattr(self, name, args[i])
    InnerType.__name__ = type_name
    console.locals[type_name] = InnerType

def create_method(console, inp):
    """Creates a method for a struct type (syntax: %method TypeName method_name ,param1 ,param2 ... return_value)"""
    words = inp.strip().split(' ', 2)
    if len(words) <= 2:
        console.error("Invalid amount of arguments")
        return
    type_name, method_name, rest = words
    if not type_name.isidentifier() or type_name not in console.locals:
        console.error("Invalid type name")
        return
    elif not method_name.isidentifier():
        console.error("Invalid method name")
        return
    parameters = []
    while True:
        ident, _rest = split_by_sep(rest, ' ')
        if not ident or ident[0] != ',' or not ident[1:].isidentifier():
            break
        rest = _rest
        parameters.append(ident[1:])
    if len(parameters) == 0:
        console.error("Invalid amount of parameters")
        return
    def inner(*args):
        if len(args) != len(parameters):
            raise TypeError(f"Invalid amount of arguments passed to struct method {method_name}")
        env = {}
        for i, name in enumerate(parameters):
            env[name] = args[i]
        env = {**console.locals, **env}
        return eval(rest, {}, env)
    inner.__name__ = method_name
    setattr(console.locals[type_name], method_name, inner)

def load_file(console, inp):
    """Runs file in the context of the interpreter. Commands such as these work as well."""
    file_name = inp.strip()
    try:
        f = open(file_name)
        lines = f.readlines()
        f.close()
    except:
        console.error(f"Invalid filename '{inp}'")
        return
    console.file_name = file_name
    for n, l in enumerate(lines):
        console.file_line = n + 1
        console.push(console.handle(l.strip()))
        if console.error_flag:
            console.error_flag = False
            console.file_line = -1
            console.file_name = ""
            return
    console.file_line = -1
    console.file_name = ""

commands = {
    "help" : show_help,
    "ptoggle" : ptoggle,
    "inspect" : inspect,
    "methods" : methods,
    "type" : create_type,
    "method" : create_method,
    "load" : load_file,
}

console = CustomConsole(commands)
console.locals["__console__"] = console
# For convenience sake
console.locals["sqrt"] = sqrt
console.locals["sin"] = sin
console.locals["cos"] = cos
console.locals["tan"] = tan
console.locals["log"] = log
console.locals["fact"] = factorial
console.locals["gcd"] = gcd
console.locals["randint"] = randint
console.locals["factors"] = factors
console.locals["e"] = e
console.locals["pi"] = pi
console.interact()