rewrote family tree and person classes

family_tree.py

 """Family Tree project for CSCI 30."""
 
+"""
+TODO:
+Don't initialize with root.
+Find a way to get root afterwards.
+Find a way to get isolated nodes.
+
+current structure:
+    bin tree inside a dict
+
+
+Problems to consider --
+What to do if not rooted?
+Multiple roots?
+What if the file describes two separate trees at once?
+Handling value errors.
+Incest and cycles:
+    > Adding someone already in the tree as a child
+
+Problems taken care of(?) --
+    Adding child to someone who already has two children.
+
+Questions:
+How does a dict compare object keys?
+How to deal with multiple trees?
+Is there only one root?
+
+Total nodes of the tree?
+11 generations = 2^11 - 1
+
+
+KEY = NAME????
+MAKE THE KEY ONLY THE PERSON'S NAME
+MAKE HELPER FUNC
+REWRITE EVERYTHING
+
+
+How do class variables work?
+"""
+
 
 class FamilyTree:
     """A family tree.
 
     Family Tree dict format --
-        key = person
-        value = [name of child1, name of child2, name of parent]
+        key = name of the person
+        value = a Person object
     """
 
-    """
-    TODO:
-    Don't initialize with root.
-    Find a way to get root afterwards.
-    Find a way to get isolated nodes.
-
-    Problems to consider --
-    What to do if not rooted?
-    Multiple roots?
-    What if the file describes two separate trees at once?
-    Handling value errors.
-    Incest and cycles:
-        > Adding someone already in the tree as a child
+    def __init__(self):
+        """Initialize the family tree."""
+        self.family_tree = dict()
+        self.root = None
 
-    Problems taken care of(?) --
-        Adding child to someone who already has two children.
+    def find_root(self):
+        """Find the root of the tree.
 
-    Questions:
-    How does a dict compare object keys?
-    How to deal with multiple trees?
-    Is there only one root?
+        note: work in progress
+        Start from arbitrary node,
+        then move upwards. (doesn't check for multiple roots)
 
+        or ?????
+
+        go through every object in the dictionary
+        and check if it has a parent (slow)
         """
+        return
 
-    def __init__(self, root):
-        """Initialize the family tree.
+    def set_root(self, person):
+        """Set the root node."""
+        self.root = person.get_name()
 
-        root: the name of the first generation person.
-        """
-        self.family_tree = dict()
-        self.root = root
-        self.family_tree[root] = [None, None, None]
-
-    def insert_person(self, person, parent):
-        """Add a person as a child to some parent."""
-        try:
-            parent_list = self.family_tree[parent]
-        except KeyError:
-            raise KeyError("Error: The parent is not in the tree.")
-        if parent_list[1] is None:
-            parent_list[1] = person
-            parent_list[person] = [None, None, parent]
-        elif parent_list[2] is None:
-            parent_list[2] = person
-            parent_list[person] = [None, None, parent]
-        else:
-            raise ValueError("Error: That parent already has two children.")
+    def insert_person(self, person):
+        """Add a Person object to the dict with the name as the key."""
+        self.family_tree[person.get_name()] = person
+
+    def get_person(self, name):
+        """Get the named Person object from the dict."""
+        return self.family_tree[name]
 
     def get_ancestors(self, name):
-        """Return a list of all the ancestors of a person."""
+        """Return a sorted list of the person's ancestors."""
+        person = self.family_tree[name]
+        parent = person.get_parent()
         ancestors = []
-        parent = self.family_tree[name][2]
         while parent is not None:
-            ancestors += parent
-            parent = self.family_tree[parent][2]
-        ancestors.sort()
+            ancestors.append(parent.get_name())
+            parent = parent.get_parent()
         return ancestors
 
     def get_descendants(self, name):
-        """Return a list of a person's descendants."""
+        """Return a sorted list of a person's descendants."""
+        def get_preorder(person):
+            """Get the inorder traversal given a person as the root."""
+            res = []
+            if person is not None:
+                lc, rc = person.get_children()
+                res.append(person.get_name())
+                res += get_preorder(lc)
+                res += get_preorder(rc)
+            return res
+
+        descendants = get_preorder(self.family_tree[name])[1:]
+        descendants.sort()
+        return descendants
+
+    def __str__(self):
+        """Return the name."""
+        return self.name
 
 
 class Person:
-    """A person.
+    """The nodes of the family tree.
 
     Attributes --
-        Name: a string indicating the name
-        Age: a integer indicating the age
-        Sex: a boolean, male = True, female = False
+        name: a string indicating the name
+        age: an integer indicating the age
+        sex: a boolean; male = True, female = False
+        children: a list of the children
+        parent: the parent
     """
 
     def __init__(self, name, age, sex):
-        """Initialize the Person class."""
+        """Initialize the Person object."""
         self.name = name
         self.age = age
         self.sex = sex
+        self.children = [None, None]
+        self.parent = None
+
+    def set_parent(self, parent):
+        """Set parent.
+
+        Also sets the current node as the parent's child.
+        """
+        self.parent = parent
+        self.parent.add_child(self)
+
+    def get_parent(self):
+        """Get parent."""
+        return self.parent
+
+    def add_child(self, child):
+        """Add a child.
+
+        If there are already 2 children, raises error.
+        """
+        if self.children[0] is None:
+            self.children[0] = child
+        elif self.children[1] is None:
+            self.children[1] = child
+        else:
+            raise Exception("This node already has two children.")
+
+    def get_children(self):
+        """Return a list of the children."""
+        return self.children
 
     def set_name(self, name):
         """Set name."""
@@ -109,3 +180,18 @@ class Person:
     def get_sex(self):
         """Get sex (bool)."""
         return self.sex
+
+
+if __name__ == "__main__":
+    ftree = FamilyTree()
+    A = Person("A", 69, True)
+    B = Person("B", 69, True)
+    C = Person("C", 69, True)
+    B.set_parent(A)
+    C.set_parent(A)
+    ftree.insert_person(A)
+    ftree.insert_person(B)
+    ftree.insert_person(C)
+    ftree.set_root(A)
+    print(ftree.get_ancestors(C.get_name()))
+    print(ftree.get_descendants(A.get_name()))