Quick and dirty animated svg of cluster upgrade

imagegen-2.py

import svgwrite
import random
import copy


containerSide = 14
innerMargin = 2


class Node(object):
    """Any node."""
    replaceTimeHalf = 3

    def __init__(self, x, y, width, height, isNew):
        """Init."""
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.isNew = isNew
        self.containerXYs = []
        self.containers = [None]*4
        for row in [0, 1]:
            for column in [0, 1]:
                xpos = x + column * (innerMargin + containerSide) + innerMargin
                ypos = y + row * (innerMargin + containerSide) + innerMargin
                self.containerXYs.append((xpos, ypos))

    def addContainer(self, container):
        self.containers[self.containers.index(None)] = container

    def removeContainer(self, container):
        self.containers[self.containers.index(container)] = None

    def containerLen(self):
        return len([c for c in self.containers if c])

    def addToImage(self, drawing):
        """Add to drawing."""
        self.drawing = drawing
        self.rect = self.drawing.add(svgwrite.shapes.Rect(
            insert=(self.x, self.y),
            size=(self.width, self.height),
            fill="orange" if self.isNew else "white",
            opacity="0" if self.isNew else "1",
            stroke="black",
            **{"stroke-width": "2"}
        ))
        if self.isNew:
            self.rect.add(self.drawing.animate(
                attributeName='opacity',
                begin="1s",
                to="1",
                dur="{}s".format(self.replaceTimeHalf),
                fill="freeze",
                **{"from": "0"}
            ))
        return self.rect

    def stop(self, startTime):
        begin = "{}s".format(startTime)
        self.rect.add(self.drawing.animate(
            attributeName='opacity',
            begin=begin,
            to="0",
            dur="{}s".format(self.replaceTimeHalf),
            fill="freeze",
            **{"from": "1"}
        ))

    def start(self, startTime):
        begin = "{}s".format(startTime)
        self.rect.add(self.drawing.animate(
            attributeName='opacity',
            begin=begin,
            to="1",
            dur="{}s".format(self.replaceTimeHalf),
            fill="freeze",
            **{"from": "0"}
        ))
        self.rect.add(self.drawing.animate(
            attributeName='fill',
            begin=begin,
            to="orange",
            fill="freeze",
            dur="1s",
            **{"from": "white"}
        ))


class Container(object):
    """Container."""
    drainDuration = "1s"

    def __init__(self, node, index):
        """Init."""
        self.node = node

    def index(self):
        return self.node.containers.index(self)

    def getXYInNode(self):
        """Get XY in node."""
        return self.node.containerXYs[self.index()]

    def addToImage(self, drawing):
        """Add to drawing."""
        self.drawing = drawing
        self.origXY = self.node.containerXYs[self.index()]
        self.rect = drawing.add(svgwrite.shapes.Rect(
            insert=self.origXY,
            size=(containerSide, containerSide),
            fill="blue"
        ))

    def move(self, startTime, newNode):
        begin = "{}s".format(startTime)
        newNode.addContainer(self)
        self.node.removeContainer(self)
        self.node = newNode
        newXY = self.getXYInNode()
        self.rect.add(self.drawing.animateTransform(
            'translate',
            'transform',
            dur=self.drainDuration,
            to="{} {}".format(newXY[0]-self.origXY[0], newXY[1]-self.origXY[1]),
            begin=begin,
            fill="freeze"
        ))


def drawBasic(nodes, containers):
    """Render."""
    dwg = svgwrite.Drawing('basic.svg', size=("200", "250"))
    for node in nodes:
        node.addToImage(dwg)
    for container in containers:
        container.addToImage(dwg)
    startTime = 1
    for node in nodes:
        drain(dwg, startTime, node, nodes, containers)
        startTime += 1
        node.stop(startTime)
        startTime += Node.replaceTimeHalf
        node.start(startTime)
        startTime += Node.replaceTimeHalf
    dwg.save()


def drawDoubleHalf(nodes, doubledNodes, containers):
    """Render."""
    dwg = svgwrite.Drawing('doublehalf.svg', size=("350", "250"))
    for node in nodes:
        node.addToImage(dwg)
    for node in doubledNodes:
        node.addToImage(dwg)

    for container in containers:
        container.addToImage(dwg)
    startTime = Node.replaceTimeHalf + 1
    for node in nodes:
        drain(dwg, startTime, node, doubledNodes, containers)
        startTime += 1
    for node in nodes:
        node.stop(startTime)
    dwg.save()


def drain(dwg, begin, node, nodes, containers):
    """Drain."""
    toBeMoved = [c for c in containers if c.node == node]
    for container in toBeMoved:
        minFill = min([n.containerLen() for n in nodes if n != node])
        choices = [n for n in nodes if n.containerLen() == minFill and n != node]
        newNode = random.choice(choices)
        container.move(begin, newNode)


def main():
    """Main."""
    nodeColumns = 4
    nodeRows = 5
    nodes = []
    containers = []
    doubledNodes = []
    margin = 5
    nodeSide = innerMargin*3 + containerSide*2

    for row in range(nodeRows):
        for column in range(nodeColumns):
            xpos = column * (margin + nodeSide) + margin
            ypos = row * (margin + nodeSide) + margin
            newNode = Node(xpos, ypos,
                           nodeSide, nodeSide, False)
            nodes.append(newNode)
            for i in xrange(random.randint(1, 4)):
                container = Container(newNode, i)
                containers.append(container)
                newNode.addContainer(container)
            xpos2 = xpos + 4*margin + nodeColumns * (margin + nodeSide)
            newDoubledNode = Node(xpos2, ypos,
                                  nodeSide, nodeSide, True)
            doubledNodes.append(newDoubledNode)
    basicNodes = copy.deepcopy(nodes)
    basicContainers = [c for n in basicNodes for c in n.containers if c]
    drawBasic(basicNodes, basicContainers)
    drawDoubleHalf(nodes, doubledNodes, containers)


if __name__ == '__main__':
    main()