Factor tree using graphviz

factorize.py

#!/usr/bin/env python3
'''
Simple script which creates a factor tree in pdf format for any given
number using the graphviz library. Did on a whim because I've never used
graphviz before. Maybe should have done MTBDDs rather than factor trees?
Oh, well, maybe next time.

No warranty is provided. MIT license.

By ifndefJOSH on GitHub
'''

import graphviz
import uuid

def isPrime(n):
	for i in range(2, n):
		if n % i == 0:
			return False
	return True

# Use global primes list because creating list of primes is expensive
primes = []

def factorize(n):
	global primes
	# Initialize primes
	topPrimeCandidate = 2 if len(primes) == 0 else primes[len(primes) - 1]
	# only create primes that could be factors of n
	while topPrimeCandidate <= n:
		if isPrime(topPrimeCandidate):
			primes.append(topPrimeCandidate)
		topPrimeCandidate += 1

	# Yeah this doesn't create a balanced tree, but whatever
	for prime in primes:
		if n == prime:
			return prime
		elif n % prime == 0:
			arr = factorize(n // prime)
			return {n:[arr, prime]}
	# Shouldn't get here unless initialization is wrong
	# but this is to ensure it is correct even if it did.
	return n

def createGraph(factorTree, dot=None, level=0, grandparent=None):
	# dotWasNone = dot is None
	if dot is None:
		dot = graphviz.Digraph(comment="Factor Tree")

	# Hackey graphviz code
	if type(factorTree) is int:
		nodeId = str(uuid.uuid4())
		dot.node(nodeId, f"{factorTree}")
		if grandparent is not None:
			dot.edge(grandparent, nodeId, constraint='true')
	elif type(factorTree) is dict:
		for parent, childList in factorTree.items():
			parentId = str(uuid.uuid4())
			dot.node(parentId, f"{parent}")
			if grandparent is not None:
				dot.edge(grandparent, parentId, constraint='true')
			for child in childList:
				createGraph(child, dot, level + 1, parentId)
	else:
		print(f"Warning, unknown type {type(factorTree)}")

	if grandparent is None:
		dot.render("factor-tree.gv", view=True)

if __name__=="__main__":
	n = int(input("Please put a number to factorize: "))
	graph = factorize(n)
	print(graph)
	createGraph(graph)