Mandelbrot ASCII art from PyPy (independent version)

ansi_mandelbrot.py

import os
import sys
import colorsys

"""
Black       0;30     Dark Gray     1;30
Blue        0;34     Light Blue    1;34
Green       0;32     Light Green   1;32
Cyan        0;36     Light Cyan    1;36
Red         0;31     Light Red     1;31
Purple      0;35     Light Purple  1;35
Brown       0;33     Yellow        1;33
Light Gray  0;37     White         1;37
"""

# used for debugging/finding new coordinates
# How to:
#   1. Set DEBUG to True
#   2. Add a new coordinate to coordinates with a high distance and high max colour (e.g. 300)
#   3. Run, pick an interesting coordinate from the shown list and replace the newly added
#      coordinate by it.
#   4. Rerun to see the max colour, insert this max colour where you put the high max colour.
#   5. Set DEBUG to False
DEBUG = False


try:
    from io import StringIO
except ImportError:
    from StringIO import StringIO

try:
    callable = callable
except NameError:
    def callable(obj):
        return hasattr(obj, "__call__")

if sys.version_info >= (3, 0):
    text = str
    bytes = bytes
    TextIO = StringIO
    exec ("print_ = print")
else:
    text = unicode
    bytes = str
    next = lambda it: it.next()
    class TextIO(StringIO):
        def write(self, data):
            if not isinstance(data, unicode):
                data = unicode(data, getattr(self, '_encoding', 'UTF-8'), 'replace')
            StringIO.write(self, data)

    def print_(*args, **kwargs):
        """ minimal backport of py3k print statement. """
        sep = ' '
        if 'sep' in kwargs:
            sep = kwargs.pop('sep')
        end = '\n'
        if 'end' in kwargs:
            end = kwargs.pop('end')
        file = 'file' in kwargs and kwargs.pop('file') or sys.stdout
        if kwargs:
            args = ", ".join([str(x) for x in kwargs])
            raise TypeError("invalid keyword arguments: %s" % args)
        at_start = True
        for x in args:
            if not at_start:
                file.write(sep)
            file.write(str(x))
            at_start = False
        file.write(end)


win32_and_ctypes = False
if sys.platform == "win32":
    try:
        import ctypes
        win32_and_ctypes = True
    except ImportError:
        pass

def hsv2ansi(h, s, v):
    # h: 0..1, s/v: 0..1
    if s < 0.1:
        return int(v * 23) + 232
    r, g, b = [int(x * 5) for x in colorsys.hsv_to_rgb(h, s, v)]
    return 16 + (r * 36) + (g * 6) + b

def ramp_idx(i, num):
    assert num > 0
    i0 = float(i) / num
    h = 0.57 + i0
    s = 1 - pow(i0,3)
    v = 1
    return hsv2ansi(h, s, v)

def ansi_ramp(num):
    return [ramp_idx(i, num) for i in range(num)]

if os.environ.get('TERM', 'dumb').find('256') > 0:
    palette = ["38;5;%d" % x for x in ansi_ramp(80)]
else:
    palette = [39, 34, 35, 36, 31, 33, 32, 37]

def _getdimensions():
    import termios,fcntl,struct
    call = fcntl.ioctl(1,termios.TIOCGWINSZ,"\000"*8)
    height,width = struct.unpack( "hhhh", call ) [:2]
    return height, width


def get_terminal_width():
    try:
        height, width = _getdimensions()
    except (KeyboardInterrupt, SystemExit, MemoryError, GeneratorExit):
        raise
    except:
        # FALLBACK
        width = int(os.environ.get('COLUMNS', 80))
    else:
        # XXX the windows getdimensions may be bogus, let's sanify a bit
        if width < 40:
            width = 80
    return width

# XXX unify with _escaped func below
def ansi_print(text, esc, file=None, newline=True, flush=False):
    if file is None:
        file = sys.stderr
    text = text.rstrip()
    if esc and not isinstance(esc, tuple):
        esc = (esc,)
    if esc and sys.platform != "win32" and file.isatty():
        text = (''.join(['\x1b[%sm' % cod for cod in esc])  +
                text +
                '\x1b[0m')     # ANSI color code "reset"
    if newline:
        text += '\n'

    if esc and win32_and_ctypes and file.isatty():
        if 1 in esc:
            bold = True
            esc = tuple([x for x in esc if x != 1])
        else:
            bold = False
        esctable = {()   : FOREGROUND_WHITE,                 # normal
                    (31,): FOREGROUND_RED,                   # red
                    (32,): FOREGROUND_GREEN,                 # green
                    (33,): FOREGROUND_GREEN|FOREGROUND_RED,  # yellow
                    (34,): FOREGROUND_BLUE,                  # blue
                    (35,): FOREGROUND_BLUE|FOREGROUND_RED,   # purple
                    (36,): FOREGROUND_BLUE|FOREGROUND_GREEN, # cyan
                    (37,): FOREGROUND_WHITE,                 # white
                    (39,): FOREGROUND_WHITE,                 # reset
                    }
        attr = esctable.get(esc, FOREGROUND_WHITE)
        if bold:
            attr |= FOREGROUND_INTENSITY
        STD_OUTPUT_HANDLE = -11
        STD_ERROR_HANDLE = -12
        if file is sys.stderr:
            handle = GetStdHandle(STD_ERROR_HANDLE)
        else:
            handle = GetStdHandle(STD_OUTPUT_HANDLE)
        oldcolors = GetConsoleInfo(handle).wAttributes
        attr |= (oldcolors & 0x0f0)
        SetConsoleTextAttribute(handle, attr)
        while len(text) > 32768:
            file.write(text[:32768])
            text = text[32768:]
        if text:
            file.write(text)
        SetConsoleTextAttribute(handle, oldcolors)
    else:
        file.write(text)

    if flush:
        file.flush()

def should_do_markup(file):
    if os.environ.get('PY_COLORS') == '1':
        return True
    if os.environ.get('PY_COLORS') == '0':
        return False
    return hasattr(file, 'isatty') and file.isatty() \
           and os.environ.get('TERM') != 'dumb' \
           and not (sys.platform.startswith('java') and os._name == 'nt')


if win32_and_ctypes:
    import ctypes
    from ctypes import wintypes

    # ctypes access to the Windows console
    STD_OUTPUT_HANDLE = -11
    STD_ERROR_HANDLE  = -12
    FOREGROUND_BLACK     = 0x0000 # black text
    FOREGROUND_BLUE      = 0x0001 # text color contains blue.
    FOREGROUND_GREEN     = 0x0002 # text color contains green.
    FOREGROUND_RED       = 0x0004 # text color contains red.
    FOREGROUND_WHITE     = 0x0007
    FOREGROUND_INTENSITY = 0x0008 # text color is intensified.
    BACKGROUND_BLACK     = 0x0000 # background color black
    BACKGROUND_BLUE      = 0x0010 # background color contains blue.
    BACKGROUND_GREEN     = 0x0020 # background color contains green.
    BACKGROUND_RED       = 0x0040 # background color contains red.
    BACKGROUND_WHITE     = 0x0070
    BACKGROUND_INTENSITY = 0x0080 # background color is intensified.

    SHORT = ctypes.c_short
    class COORD(ctypes.Structure):
        _fields_ = [('X', SHORT),
                    ('Y', SHORT)]
    class SMALL_RECT(ctypes.Structure):
        _fields_ = [('Left', SHORT),
                    ('Top', SHORT),
                    ('Right', SHORT),
                    ('Bottom', SHORT)]
    class CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure):
        _fields_ = [('dwSize', COORD),
                    ('dwCursorPosition', COORD),
                    ('wAttributes', wintypes.WORD),
                    ('srWindow', SMALL_RECT),
                    ('dwMaximumWindowSize', COORD)]

    _GetStdHandle = ctypes.windll.kernel32.GetStdHandle
    _GetStdHandle.argtypes = [wintypes.DWORD]
    _GetStdHandle.restype = wintypes.HANDLE
    def GetStdHandle(kind):
        return _GetStdHandle(kind)

    SetConsoleTextAttribute = ctypes.windll.kernel32.SetConsoleTextAttribute
    SetConsoleTextAttribute.argtypes = [wintypes.HANDLE, wintypes.WORD]
    SetConsoleTextAttribute.restype = wintypes.BOOL

    _GetConsoleScreenBufferInfo = \
        ctypes.windll.kernel32.GetConsoleScreenBufferInfo
    _GetConsoleScreenBufferInfo.argtypes = [wintypes.HANDLE,
                                ctypes.POINTER(CONSOLE_SCREEN_BUFFER_INFO)]
    _GetConsoleScreenBufferInfo.restype = wintypes.BOOL
    def GetConsoleInfo(handle):
        info = CONSOLE_SCREEN_BUFFER_INFO()
        _GetConsoleScreenBufferInfo(handle, ctypes.byref(info))
        return info

    def _getdimensions():
        handle = GetStdHandle(STD_OUTPUT_HANDLE)
        info = GetConsoleInfo(handle)
        # Substract one from the width, otherwise the cursor wraps
        # and the ending \n causes an empty line to display.
        return info.dwSize.Y, info.dwSize.X - 1


class Mandelbrot:
    def __init__ (self, width=100, height=28, x_pos=-0.5, y_pos=0, distance=6.75):
        self.xpos = x_pos
        self.ypos = y_pos
        aspect_ratio = 1/3.
        factor = float(distance) / width # lowering the distance will zoom in
        self.xscale = factor * aspect_ratio
        self.yscale = factor
        self.iterations = 170
        self.x = width
        self.y = height
        self.z0 = complex(0, 0)

    def init(self):
        self.reset_lines = False
        xmin = self.xpos - self.xscale * self.x / 2
        ymin = self.ypos - self.yscale * self.y / 2
        self.x_range = [xmin + self.xscale * ix for ix in range(self.x)]
        self.y_range = [ymin + self.yscale * iy for iy in range(self.y)]

    def reset(self, cnt):
        self.reset_lines = cnt

    def generate(self):
        self.reset_lines = False
        iy = 0
        while iy < self.y:
            ix = 0
            while ix < self.x:
                c = complex(self.x_range[ix], self.y_range[iy])
                z = self.z0
                colour = 0
                mind = 2

                for i in range(self.iterations):
                    z = z * z + c
                    d = abs(z)
                    if d >= 2:
                        colour = min(int(mind / 0.007), 254) + 1
                        break
                    else:
                        mind = min(d, mind)

                yield ix, iy, colour
                if self.reset_lines is not False: # jump to the beginning of the line
                    iy += self.reset_lines
                    do_break = bool(self.reset_lines)
                    self.reset_lines = False
                    if do_break:
                        break
                    ix = 0
                else:
                    ix += 1
            iy += 1


class Driver(object):
    zoom_locations = [
        # x, y, "distance", max color range
        (0.37865401, 0.669227668, 0.04, 111),
        (-1.2693, -0.4145, 0.2, 105),
        (-1.2693, -0.4145, 0.05, 97),
        (-1.2642, -0.4185, 0.01, 95),
        (-1.15, -0.28, 0.9, 94),
        (-1.15, -0.28, 0.3, 58),
        (-1.15, -0.28, 0.05, 26),
            ]
    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self.zoom_location = -1
        self.max_colour = 256
        self.colour_range = None
        self.invert = True
        self.interesting_coordinates = []
        self.init()

    def init(self):
        self.width = get_terminal_width() or 80 # in some envs, the py lib doesnt default the width correctly
        self.mandelbrot = Mandelbrot(width=(self.width or 1), **self.kwargs)
        self.mandelbrot.init()
        self.gen = self.mandelbrot.generate()

    def reset(self, cnt=0):
        """ Resets to the beginning of the line and drops cnt lines internally. """
        self.mandelbrot.reset(cnt)

    def restart(self):
        """ Restarts the current generator. """
        print_(file=sys.stderr)
        self.init()

    def dot(self):
        """ Emits a colourful character. """
        x = c = 0
        try:
            x, y, c = next(self.gen)
            if x == 0:
                width = get_terminal_width()
                if width != self.width:
                    self.init()
        except StopIteration:
            if DEBUG and self.interesting_coordinates:
                print_("Interesting coordinates:", self.interesting_coordinates, file=sys.stderr)
                self.interesting_coordinates = []
            kwargs = self.kwargs
            self.zoom_location += 1
            self.zoom_location %= len(self.zoom_locations)
            loc = self.zoom_locations[self.zoom_location]
            kwargs.update({"x_pos": loc[0], "y_pos": loc[1], "distance": loc[2]})
            self.max_colour = loc[3]
            if DEBUG:
                # Only used for debugging new locations:
                print_("Colour range", self.colour_range)
            self.colour_range = None
            self.restart()
            return
        if self.print_pixel(c, self.invert):
            self.interesting_coordinates.append(dict(x=(x, self.mandelbrot.x_range[x]),
                                                     y=(y, self.mandelbrot.y_range[y])))
        if x == self.width - 1:
            print_(file=sys.stderr)

    def print_pixel(self, colour, invert=1):
        chars = [".", ".", "+", "*", "%", "#"]
        idx = lambda chars: (colour+1) * (len(chars) - 1) // self.max_colour
        if invert:
            idx = lambda chars, idx=idx:len(chars) - 1 - idx(chars)
        char = chars[idx(chars)]
        ansi_colour = palette[idx(palette)]
        ansi_print(char, ansi_colour, newline=False, flush=True)
        if DEBUG:
            if self.colour_range is None:
                self.colour_range = [colour, colour]
            else:
                old_colour_range = self.colour_range
                self.colour_range = [min(self.colour_range[0], colour), max(self.colour_range[1], colour)]
                if old_colour_range[0] - colour > 3 or colour - old_colour_range[1] > 3:
                    return True


if __name__ == '__main__':
    import random
    from time import sleep

    d = Driver()
    while True:
        sleep(random.random() / 800)
        d.dot()
        if 0 and random.random() < 0.01:
            string = "WARNING! " * 3
            d.jump(len(string))
            print_(string, end=' ')