Python fractal tree using SVG

Question

I made a simple program to generate fractal tree using PATH object of SVG. Any suggestions?

import math
from random import randint

# const - upper limit for randint
s = 10

"""Simple fractal tree using SVG and recursion.

Usage:
    Create Root object bt Branch(x1=400, y1=800, x2=400, y2=600, color=60, size=35)
    x1, y1, x2, y2 - start points of root

    Generate Tree Tree(lenght=200, angle=-20, depth=9, x1=400, y1=600, size=35, color=60, outlist=resutlist)

    lenght - lenght of start branch
    angle - start angle of branch
    depth - number of tree level
    x1, y1 - start point of branch
"""


class Branch():
    """Class represents a single branch."""
    def __init__(self, x1, y1, x2, y2, color, size):
        """Assigning  values."""
        self.x1 = x1
        self.x2 = x2
        self.y1 = y1
        self.y2 = y2
        self.color = color
        self.size = size

    def __str__(self):
        """Return path SVG object with points, color and stroke of branch."""
        return '<path d="M {x1} {y1} L {x2} {y2}" stroke="rgb(100,{c},0)" stroke-width="{w}"/>\n'.format(
            x1=self.x1,
            y1=self.y1,
            x2=self.x2,
            y2=self.y2,
            w=self.size,
            c=self.color
        )

    def __repr__(self):
        """Return text represent object."""
        return self.__str__()


class Tree():
    """
    Class represents Tree.

    Tree is composed of Branch object.
    """
    def __init__(self, lenght, angle, depth, x1, y1, size, color, outlist):
        """Main point of start generation."""
        self.branches = self.drawbranch(lenght, angle, depth, x1, y1, size, color, outlist)

    def drawbranch(self, lenght, angle, depth, x1, y1, size, color, outlist):
        """Recursive function for generate three Branch object per iteration."""
        # if depth > 0
        if depth:
            # X value of second point
            x2 = x1 + lenght * math.cos(math.radians(angle))
            # Y value of second point
            y2 = y1 + lenght * math.sin(math.radians(angle))

            # modify lenght of single branch
            lenght = float(2.0 / 3.0 * lenght)
            # modify size of single branch
            size = float(2.0 / 3.0 * size) + 1
            # modify color of single branch
            color += 6

            # X value of B point
            bx = x1
            # Y value of B point
            by = y2

            # X value of C point
            cx = -x2 + 2 * x1
            # Y value of C point
            cy = y2

            # Create A point
            b1 = Branch(x1, y1, x2, y2, color, size)
            # Add to list
            outlist.append(str(b1))
            # Call drawbranch function (recursion)
            self.drawbranch(lenght, angle + randint(-10, s), depth - 1, x2, y2, size, color, outlist)

            # Create B point
            b2 = Branch(x1, y1, bx, by, color, size)
            # Add to list
            outlist.append(str(b2))
            # Calculate new angle
            nangle = angle + randint(-1, 0) * randint(1, s)
            # Call drawbranch function (recursion)
            self.drawbranch(lenght, nangle, depth - 1, bx, by, size, color, outlist)

            # Create C point
            b3 = Branch(x1, y1, cx, cy, color, size)
            # Add to list
            outlist.append(str(b3))
            # Calculate new angle
            nangle = angle + randint(0, 1) * randint(1, s)
            # Call drawbranch function (recursion)
            self.drawbranch(lenght, nangle, depth - 1, cx, cy, size, color, outlist)

        # Return list of branches
        return outlist

    def write_svg(self, output='drzewko.svg'):
        """Function that write all branches to SVG file."""
        with open(output, 'w') as outfile:
            # Write SVG declaration
            outfile.write('<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 800 800" version="1.1">\n')
            # Map to str all branches and write it into file
            outfile.writelines(map(str, self.branches))
            # End of SVG file
            outfile.write('</svg>\n')


print "Start generating, please wait.."
# Create empty list
resutlist = []
# Create root of Tree and add to list
resutlist.append(Branch(x1=400, y1=800, x2=400, y2=600, color=60, size=35))

# Call Tree object
t = Tree(lenght=200, angle=-20, depth=9, x1=400, y1=600, size=35, color=60, outlist=resutlist)
# After generate tree save to file
t.write_svg()

print "Done, check SVG file"

Results:

Answer 1

Interesting problem thanks for sharing.

I have refactored your code and it is included in it entirety below. The most significant change made was to introduce a Point class, to encapsulate (x, y) pairs, and to provide a bit of syntactic sugar to math on same.

Make Point as class:

So this code:

# X value of C point
cx = -x2 + 2 * x1
# Y value of C point
cy = y2

can become:

c = Point(-p2.x + 2 * p1.x, p2.y)

via a namedtuple that provides an immutable datastructure whose attributes can be accessed with things like p2.x.

Comments:

I removed many of your comments. In several of these cases, I renamed a variable to include the information that was present in the comment.

Use object attributes

Before the edits, the variable branches was being passed recursively, but it was eventually assigned as an attribute of the Tree class. So the new code just starts there, and branches is an attribute of the class from the beginning.

Type cast to float?

In general you do not need to cast floats to floats. So something like:

 branch_length = float(2.0 / 3.0 * length)

can just be:

 branch_length = 2.0 / 3.0 * length

since 2.0 is already a float.

Complete Code:

import math
from random import randint
from collections import namedtuple

# const - upper limit for randint
s = 10


class Point(namedtuple('Point', 'x y')):

    def __str__(self):
        return'{} {}'.format(self.x, self.y)

    def __add__(self, other):
        assert isinstance(other, Point)
        return Point(self.x + other.x, self.y + other.y)

    def __mul__(self, other):
        return Point(self.x * other, self.y * other)

    def __rmul__(self, other):
        return self.__mul__(other)


class Branch(namedtuple('Branch', 'p1 p2 color size')):

    def __str__(self):
        """Path SVG object with points, color and stroke of branch."""
        return ('<path d="M {p1} L {p2}" '
                'stroke="rgb(100,{c},0)" stroke-width="{w}"/>\n'.
                format(p1=self.p1, p2=self.p2, w=self.size, c=self.color))

    def __repr__(self):
        return self.__str__()


class Tree(object):

    def __init__(self, length, angle, depth, point, size, color, outlist):
        """Main point of start generation."""
        self.branches = outlist
        self.draw_branches(length, angle, depth, point, size, color)

    def draw_branches(self, length, angle, depth, p1, size, color):
        """ Recursively generate three Branch objects per iteration."""
        if depth <= 0:
            return

        p2 = p1 + length * Point(
            math.cos(math.radians(angle)),
            math.sin(math.radians(angle))
        )

        # set some new characteristics for the next level
        branch_length = 2.0 / 3.0 * length
        branch_size = 2.0 / 3.0 * size + 1
        color += 6

        # Calculate new angle and recurse
        self.branches.append(Branch(p1, p2, color, branch_size))
        nangle = angle + randint(-10, s)
        self.draw_branches(branch_length, nangle, depth - 1,
                           p2, branch_size, color)

        # Calculate new angle and recurse
        b = Point(p1.x, p2.y)
        self.branches.append(Branch(p1, b, color, branch_size))
        nangle = angle + randint(-1, 0) * randint(1, s)
        self.draw_branches(branch_length, nangle, depth - 1,
                           b, branch_size, color)

        # Calculate new angle and recurse
        c = Point(-p2.x + 2 * p1.x, p2.y)
        self.branches.append(Branch(p1, c, color, branch_size))
        nangle = angle + randint(0, 1) * randint(1, s)
        self.draw_branches(branch_length, nangle, depth - 1,
                           c, branch_size, color)

    def write_svg(self, output='drzewko.svg'):
        with open(output, 'w') as outfile:
            outfile.write('<svg xmlns="http://www.w3.org/2000/svg" '
                          'viewBox="0 0 800 800" version="1.1">\n')
            outfile.writelines(map(str, self.branches))
            outfile.write('</svg>\n')


print("Start generating, please wait..")

# a starting point
resultlist = [Branch(Point(400, 800), Point(400, 600), color=60, size=35)]

# Build and save the tree as an svg
t = Tree(length=200, angle=-20, depth=9, point=Point(400, 600),
         size=35, color=60, outlist=resultlist)
t.write_svg()

print("Done, check SVG file")