How can I create a variant tree using Python?

Question

I want to assess how many variants of a given number of features and corresponding attributes exist and then plot the combinations as a tree of branches, where each branch represents one combination/variant.

Example: features: colors,size,dim. Each feature has different attributes (colors: red,green,blue; size: big, small).

Through permutation I find the number of variants. Each combination/variant is a branch of my variant tree. Eg ('red', 'big', 1) is one branch, ('red', 'big', 2) is another and so on.

Is there any libary which can help me draw out these branches with nodes and arcs?

Code for permutations:

colors = ['red', 'green', 'blue']
size = ['big','small']
dim = [1,2,3]

from itertools import product

x =list (product(colors,size,dim))

print (x)
print ("Number of variants:",len(x))

[('red', 'big', 1), ('red', 'big', 2), ('red', 'big', 3), 
 ('red', 'small', 1), ('red', 'small', 2), ('red', 'small', 3), 
 ('green', 'big', 1), ('green', 'big', 2), ('green', 'big', 3), 
 ('green', 'small', 1), ('green', 'small', 2), ('green', 'small', 3), 
 ('blue', 'big', 1), ('blue', 'big', 2), ('blue', 'big', 3), 
 ('blue', 'small', 1), ('blue', 'small', 2), ('blue', 'small', 3)]

enter image description here

Answer 1

I was able to create this using the general graph description format DOT . Python file main.py creating the DOT file:

from itertools import product

colors = ['red', 'green', 'blue']
size = ['big','small']
dim = [1,2,3]

print('digraph G {\n    rankdir=LR\n    node [shape=box]\n')
print('    start [label=""]')

for i, c in enumerate(colors):
    print(f'    color_{i} [label="{c}"]; start -> color_{i}')

    for j, s in enumerate(size):
        print(f'    size_{i}_{j} [label="{s}"]; color_{i} -> size_{i}_{j}')

        for k, d in enumerate(dim):
            print(f'    dim_{i}_{j}_{k} [label="{d}"]; size_{i}_{j} -> dim_{i}_{j}_{k}')
print('}')

I created the image using the following commands:

$ python main.py > graph.dot
$ dot -Tpng graph.dot -o graph.png

The dot command I used is from the graphviz package. I haven't used DOT much so I wasn't able to add the text and blue color.