import string
from _pyrepl.readline import multiline_input
import _pyrepl.utils

def count_brackets(s: str) -> int:
    return s.count("[") - s.count("]")

def more_lines(s: str) -> bool:
    return count_brackets(s) > 0

reader = _pyrepl.readline._wrapper.get_reader()
original_after_command = reader.after_command
def after_command(self, cmd) -> None:
    original_after_command(cmd)

    if cmd.finish: return

    buffer = "".join(self.buffer)
    msg = repr(buffer)
    if msg != self.msg:
        self.msg = msg
        self.dirty = True
reader.after_command = lambda cmd: after_command(reader, cmd) # fake instance method

if hasattr(_pyrepl.utils, 'ColorSpan'):
    from _pyrepl.utils import ColorSpan, Span
    import _pyrepl.readline

    def gen_colors(s):
        for i, c in enumerate(s):
            if c in string.ascii_lowercase:
                yield ColorSpan(Span(i, i), "string")
            elif c in string.digits or c in string.ascii_uppercase:
                yield ColorSpan(Span(i, i), "number")
            else:
                yield ColorSpan(Span(i, i), "reset")

    if hasattr(reader, 'gen_colors'):
        reader.gen_colors = gen_colors
    else:
        import _pyrepl.reader
        _pyrepl.reader.gen_colors = gen_colors

ps1 = "pydc> "
ps2 = "  ... "

while True:
    try:
        statement = multiline_input(more_lines, ps1, ps2)
    except EOFError:
        break
    except KeyboardInterrupt:
        print()
        break
    print(repr(statement))