diff options
Diffstat (limited to 'src/tree.py')
| -rw-r--r-- | src/tree.py | 148 |
1 files changed, 126 insertions, 22 deletions
diff --git a/src/tree.py b/src/tree.py index a22b3e8..f7bd330 100644 --- a/src/tree.py +++ b/src/tree.py @@ -2,6 +2,7 @@ from person import Person from union import Union import const +import pickle # Tree class @@ -15,19 +16,104 @@ class Tree: self.unions = [] - def Person(self, name, sex): + def read_from(self, file): + "Read tree from the file" + + # Load persons + count = pickle.load(file) + for i in range(count): + id = pickle.load(file) + name = pickle.load(file) + sex = pickle.load(file) + self.Person(name, sex, id) + + # Load unions + count = pickle.load(file) + for i in range(count): + # Create union + id = pickle.load(file) + dadid = pickle.load(file) + mumid = pickle.load(file) + dad = self.get_person(dadid) + if dad is None: + print "ERROR on "+str(dadid) + mum = self.get_person(mumid) + if mum is None: + print "ERROR on "+str(mumid) + u = self.Union(dad, mum, id) + + # Create childs + childcount = pickle.load(file) + for j in range(childcount): + childid = pickle.load(file) + child = self.get_person(childid) + if child is None: + print "ERROR on "+str(childid) + u.append_child(child) + + # Get root + rootid = pickle.load(file) + root = self.get_person(rootid) + if root is None: + print "ERROR on "+str(rootid) + self.root = root + + return self + + + def write_to(self, file): + "Write tree to file" + + # Save persons + pickle.dump(len(self.persons), file) + for p in self.persons: + pickle.dump(p.id, file) + pickle.dump(p.name, file) + pickle.dump(p.sex, file) + + # Save unions + pickle.dump(len(self.unions), file) + for u in self.unions: + pickle.dump(u.id, file) + pickle.dump(u.dad.id, file) + pickle.dump(u.mum.id, file) + pickle.dump(len(u.childs), file) + for c in u.childs: + pickle.dump(c.id, file) + + # Save root + pickle.dump(self.root.id, file) + + + def Person(self, name, sex, id=None): "Add a new person" - p = Person(name, sex) + p = Person(name, sex, id) self.persons.append(p) return p - def Union(self, dad, mum): + def Union(self, dad, mum, id=None): "Add a new union" - u = Union(dad, mum) + u = Union(dad, mum, id) self.unions.append(u) return u + + def get_person(self, id): + "Look for a person with an id" + for p in self.persons: + if p.id == id: + return p + return None + + + def get_union(self, id): + "Look for an union with an id" + for u in self.unions: + if u.id == id: + return u + return None + def is_descendant(self, person, root=None): "Look if person is a descendant of the tree" @@ -105,7 +191,7 @@ class Tree: root.compute_desc_pos(x, y) size = root.size_desc() if i > 0: - size = size + roots[i].size_desc() + size = size + roots[i].child_count() x = x + (size*const._person_width+size*root.width_margin())/2 x = initx y = y + const._person_height+root.height_margin() @@ -183,23 +269,9 @@ def empty_tree(): 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 +def sample_family1(): + # Large family sample tree = Tree() l = tree.Person("Lucien", "M") @@ -247,6 +319,38 @@ def sample_family2(): uc2.append_child(pi) uc2.append_child(tree.Person("Camille", "F")) - tree.root = d #vm + tree.root = d return tree + + +def sample_family2(): + # Napoleon family + tree = Tree() + tree.root = p1 = tree.Person("Chales-Marie Bonaparte", "M") + p2 = tree.Person("Letizia Ramolino", "F") + u1 = tree.Union(p2, p1) + p4 = tree.Person("Napoleon", "M") + p6 = tree.Person("Louis Bonaparte", "M") + u1.append_child(p4) + u1.append_child(p6) + p3 = tree.Person("Josephine", "F") + u2 = tree.Union(p4, p3) + p5 = tree.Person("Marie-Louise d'Autriche", "F") + u3 = tree.Union(p4, p5) + + p10 = tree.Person("Napoleon II", "M") + u3.append_child(p10) + + p7 = tree.Person("Hortense de Beauharnais", "F") + u4 = tree.Union(p6, p7) + p8 = tree.Person("Napoleon III", "M") + u4.append_child(p8) + + p9 = tree.Person("Eugenie de Palafox-Guzman", "F") + u5 = tree.Union(p8, p9) + p11 = tree.Person("Louis Napoleon", "M") + u5.append_child(p11) + + return tree + |