path: root/src/tree.py

from person import Person
from union import Union
import const


# Tree class
class Tree:
	"Class to represent a tree: a sum of persons and unions"

	def __init__(self):
		"Constructor, init fields to None"
		self.root = None
		self.persons = []
		self.unions = []
	
	
	def Person(self, name, sex):
		"Add a new person"
		p = Person(name, sex)
		self.persons.append(p)
		return p
		
	
	def Union(self, dad, mum):
		"Add a new union"
		u = Union(dad, mum)
		self.unions.append(u)
		return u
		
		
	def is_descendant(self, person, root=None):
		"Look if person is a descendant of the tree"
		start = root
		if start is None:
			start = self.root
		for u in start.unions:
			for c in u.childs:
				if c == person:
					return True
				if self.is_descendant(person, c):
					return True
		return False
		
		
	def is_ascendant(self, person, root=None):
		"Look if person is an ascendant of the tree"
		start = root
		if start is None:
			start = self.root
		if start.parents is not None:
			if start.parents.dad == person or start.parents.mum == person:
				return True
			if self.is_ascendant(person, start.parents.dad):
				return True
			if self.is_ascendant(person, start.parents.mum):
				return True
		return False
		

	def is_family(self, person, root=None):
		"Look if person is a direct family member"
		
		# Start point
		start = root
		if start is None:
			start = self.root
			
		# Test 
		if start == person:
			return True
		if self.is_descendant(person, start):
			return True
					
		# Test parents
		if start.parents is not None:
			if self.is_family(person, start.parents.dad):
				return True
			if self.is_family(person, start.parents.mum):
				return True
				
		return False	
		
		
	def set_position(self, initx, inity):
		"Set the tree to this position"
		
		# Compute levels
		levels = dict()
		maxlevel = -1
		for p in self.persons:
			levelvalue = p.level()
			maxlevel = max(levelvalue, maxlevel)
			if levelvalue not in levels:
				levels[levelvalue] = []
			levels[levelvalue].append(p)
			p.computed = False
			
		# Set pos start starting on higher level
		(x, y) = (initx, inity)
		for levelvalue in reversed(range(maxlevel+1)):
			roots = levels[levelvalue]
			for i, root in enumerate(roots):
				root.set_pos(x, y)
				root.compute_desc_pos(x, y)
				size = root.size_desc()
				if i > 0:
					size = size + roots[i].size_desc()
				x = x + (size*const._person_width+size*root.width_margin())/2
			x = initx
			y = y + const._person_height+root.height_margin()
		
	
	def draw(self, gc, pc):
		"Draw the tree in the graphical and pango context"
		for p in self.persons:
			p.draw(gc, pc)
			
		for u in self.unions:
			u.draw(gc, pc)
	
	
	def person_at(self, x, y):
		"Look for a person at this position"
		for p in self.persons:
			if p.point_inside(x, y):
				return p
		return None
		
	
	def translate(self, dx, dy):
		"Translate tree"
		for p in self.persons:
			p.translate(dx, dy)
		
		
	def scale(self, scalelevel):
		"Scale tree"
		for p in self.persons:
			p.scale(scalelevel)
			
	
	def remove(self, person):
		"Remove a person from the tree"
		# Remove from union
		for u in person.unions:
			self.unions.remove(u)
			if u.dad == person:
				u.mum.unions.remove(u)
			else:
				u.dad.unions.remove(u)

		# Remove from parent child
		if person.parents is not None:
			person.parents.childs.remove(person)
			
		# Remove from tree
		self.persons.remove(person)
		
		
	def remove_union(self, union):
		"Remove an union from the tree"
		# Remove from child
		for c in union.childs:
			c.parents = None
		
		# Remove from parent
		union.dad.unions.remove(union)
		union.mum.unions.remove(union)
		
		# Remove from tree
		self.unions.remove(union)
		
		# Remove dad and mum
		self.remove(union.dad)
		self.remove(union.mum)
			

# Create the samples family		
def empty_tree():
	tree = Tree()
	tree.root = tree.Person("", "M")
	
	return tree
	

def sample_family1():
	# Empty family
	tree = Tree()	
	tree.root = tree.Person("C", "M")
	
	a = tree.Person("A", "M")
	b = tree.Person("B", "F")
	tree.Union(a, b).append_child(tree.root)
	bb = tree.Person("BB", "F")
	tree.Union(a, bb)
	
	return tree	
	
	
def sample_family2():
	# Large family
	tree = Tree()	
	
	l = tree.Person("Lucien", "M")
	a = tree.Person("Annick", "F")
	u = tree.Union(l, a)

	d = tree.Person("Dominique", "M")
	u.append_child(d)
	au = tree.Person("Anne", "F")
	u.append_child(au)
	m = tree.Person("Madeleine", "F")
	u.append_child(m)

	jp = tree.Person("Jean-Pierre", "M")
	j = tree.Person("Julie", "F")
	up = tree.Union(jp, j)
	up.append_child(l)
	c = tree.Person("Christian", "M")
	up.append_child(c)
	
	jo = tree.Person("Jonathan", "M")
	ub = tree.Union(jo, j)
	ub.append_child(tree.Person("Charlie", "F"))
	
	rs = tree.Person("Renee", "F")
	vm = tree.Person("Vivien", "M")
	urv = tree.Union(vm, rs)
	urv.append_child(j)

	jr = tree.Person("Jean-Rene", "M")
	ua = tree.Union(jr, tree.Person("Micheline", "F"))
	ua.append_child(a)
	i = tree.Person("Irene", "F")
	ua.append_child(i)
	ui = tree.Union(tree.Person("Nathan", "M"), i)
	ui.append_child(tree.Person("Marie", "F"))
	ui.append_child(tree.Person("Noel", "M"))
	ui.append_child(tree.Person("Thierry", "M"))

	uc = tree.Union(c, tree.Person("Clarah", "F"))
	sa = tree.Person("Sandrine", "F")
	uc.append_child(sa)
	uc2 = tree.Union(c, tree.Person("Vivianne", "F"))
	pi = tree.Person("Pierre", "M")
	uc2.append_child(pi)
	uc2.append_child(tree.Person("Camille", "F"))
	
	tree.root = d #vm
	
	return tree