christmas_tree.py

import random


class ASCIICanvas:
    def __init__(self, w=50, h=100):
        assert 0 < w < 1000 and 0 < h < 1000
        self.w = w
        self.h = h
        self.buffer = []
        self.clear()
        self.pen = (0, 0)

    def clear(self):
        self.buffer = [' '] * self.w * self.h

    def __str__(self):
        w, h = self.w, self.h
        rows = []
        for i in range(0, w * h, w):
            rows.append(''.join(self.buffer[i:i + w]))
        return '\n'.join(rows)

    def __setitem__(self, key, value):
        assert type(value) == str and len(value) == 1
        x, y = key
        assert 0 <= x < self.w and 0 <= y < self.h
        self.buffer[x + self.w * y] = value

    def __getitem__(self, key):
        x, y = key
        assert 0 <= x < self.w and 0 <= y < self.h
        return self.buffer[x + self.w * y]

    def move(self, x, y):
        assert 0 <= x < self.w and 0 <= y < self.h
        self.pen = (x, y)

    def dot(self, symbol='*'):
        self[self.pen] = symbol

    def crawl(self, direction, steps, dry=False):
        dir_map = {
            'n': (0, -1, '|'),
            's': (0, 1, '|'),
            'w': (-1, 0, '-'),
            'e': (1, 0, '-'),
            'nw': (-1, -1, '\\'),
            'ne': (1, -1, '/'),
            'se': (1, 1, '\\'),
            'sw': (-1, 1, '/'),
        }
        assert direction in dir_map, 'invalid direction'
        dx, dy, symbol = dir_map[direction]
        for _ in range(steps):
            if not dry:
                self.dot(symbol)
            self.pen = self.pen[0] + dx, self.pen[1] + dy
        if not dry:
            self.dot(symbol)

    def line(self, x, y, direction, steps):
        self.move(x, y)
        self.crawl(direction, steps)


def christmas_tree(c: ASCIICanvas, n_balls_max=30, p=4, steps=6):
    c.clear()

    cx = c.w // 2

    # left side
    c.move(cx, 1)
    for s in range(steps):
        c.crawl('sw', p)
        c.crawl('e', p // 2)

    # right side
    c.move(cx + 1, 1)
    for s in range(steps):
        c.crawl('se', p)
        c.crawl('w', p - 2)

    # bottom
    c.crawl('w', steps * p + 2)

    # trunc
    trunc_h = 5
    _, y = c.pen
    c.move(cx - 2, y + 1)
    c.crawl('s', trunc_h)
    c.crawl('e', 5)
    c.crawl('n', 1, dry=True)
    c.crawl('n', trunc_h - 1)

    # balls
    y0 = 2
    y1 = steps * p + y0 - 1

    balls = ['*', 'o', '*', '@']

    def is_good_for_ball(x, y):
        neighours = [
            [x, y],
            [x, y + 1],
            [x, y - 1],
            [x + 1, y],
            [x - 1, y]
        ]
        for n in neighours:
            if c[n] != ' ':
                return False
        return True

    for y in range(y0, y1):
        a0 = y * c.w
        a1 = a0 + c.w
        line = c.buffer[a0:a1]
        start = line.index('/')
        end = line.index('\\')
        if start > 0 and end > 0:
            for _ in range(20):
                x = random.randrange(start + 1, end)
                if is_good_for_ball(x, y):
                    c[x, y] = random.choice(balls)



def frame_it(s: str):
    detected_length = s.find('\n')
    s = '|' + s.replace('\n', '|\n|') + '|'
    top_bot = '\n*' + '-' * detected_length + '*\n'
    return top_bot + s + top_bot


c = ASCIICanvas(61, 32)
christmas_tree(c)
print(frame_it(str(c)))