Python Tree/Node Class from scratch

Question

I decided to try to make my own tree style class. I may use it in a project, but it's primarily to practice.

I would appreciate any feedback for improvements, best practices, etc.

import sys

class Node(object):
    def __init__(self, name=None, data=None, parent=None):
        self.name = name
        self.data = data
        self.parent = parent
        self.children = []

    def add_child(self, child):
        # Check if Child is Node()
        if type(child) is not type(self):
            print 'Child not added.'
            print 'Child[{0}] is not type:{1}'.format(type(child),type(self))
            return
        # Check if child is self or if child is ancestor
        ancestors = self.get_ancestors()
        if child is not self and child not in ancestors:
            self.children.append(child)
            child.parent = self
        # Check Failed. Would cause infinite recursion
        else:
            print 'Child not added. Recursive loop.'
            print 'Tried to add: [{0}] to Self:[{1}] | Ancestors:{1}'.format(
                                                    child, self, ancestors)
            # Print recursion cause
            if child is self:
                print 'Child is self'
            else:
                print 'Child is ancestor'

    def get_children(self):
        return self.children

    def get_child(self, index):
        try:
            return self.children[index]
        except IndexError:
            print 'Child index not found'

    # Perhaps parent should be list. 
    # Current implementation fails if node is child of different nodes
    # only returns last parent
    def get_parent(self):
        return self.parent

    def get_ancestors(self):
        ancestor = self.parent
        ancestors = [ancestor]
        try:
            while ancestor.parent != None:
                ancestors.append(ancestor)
                parent = self.parent
                ancestor = parent.get_parent()
            return ancestors
        # Node is root. Root's parent is None
        except AttributeError:
            return []

    def get_sibblings(self):
        # Excludes self
        if self.parent is not None:
            parents_chidren = list(self.parent.children)
            parents_chidren.pop(self.get_index())
            siblings = parents_chidren
            return siblings

    def get_index(self):
        for n, child in enumerate(self.parent.children):
            if child is self:
                return n

    def remove_child(self, index):
        try:
            self.children.pop(index)
        except IndexError:
            print 'Cannot remove. Child Index not found'
        return self.children

    def print_tree(self, level=0):
        print '\t' * level + '+'*(level+1) + repr(self)
        # Check if in infinite recursion loop.
        if level > 10:
            print 'Maxium recursion exceeded.'
            raise Exception('Maximum recursion.')
            sys.exit()
        for child in self.children:
            child.print_tree(level+1)

    def __repr__(self):
        return '[Node:{0}|Parent:{1}]'.format(self.name, self.parent)

root = Node(name='root')
child1 = Node(name='child1')
root.add_child(child1)

child2 = Node(name='child2')
child1.add_child(child2)
child1.add_child(child2)
root.add_child(child2)

print '\n---TEST PRINT TREE---\n'
root.print_tree()
child1.print_tree()

print '\n---TEST RECURSION LOOP CHECKS---\n'

root.add_child(root)
child1.add_child(root)

print '\n---TEST METHODS---\n'

print 'roots children: ', root.children
print 'child1 children: ', child1.children
print 'root parent: ', root.parent
print 'child1 parent: ', child1.parent
print 'child2 parent: ', child2.parent

print 'root siblings: ', root.get_sibblings()
print 'child1 siblings: ', child1.get_sibblings()

print 'child1 index: ', child1.get_index()
print 'child2 index: ', child2.get_index()

print '\n---TEST REMOVE CHILD---\n'

print 'roots children:', root.children
print 'remove child:', root.remove_child(child2.get_index())
print 'roots children:', root.children

OUTPUT:

---TEST PRINT TREE---

+[Node:root|Parent:None]
    ++[Node:child1|Parent:[Node:root|Parent:None]]
        +++[Node:child2|Parent:[Node:root|Parent:None]]
        +++[Node:child2|Parent:[Node:root|Parent:None]]
    ++[Node:child2|Parent:[Node:root|Parent:None]]
+[Node:child1|Parent:[Node:root|Parent:None]]
    ++[Node:child2|Parent:[Node:root|Parent:None]]
    ++[Node:child2|Parent:[Node:root|Parent:None]]

---TEST RECURSION LOOP CHECKS---

Child not added. Recursive loop.
Tried to add: [[Node:root|Parent:None]] to Self:[[Node:root|Parent:None]] | Ancestors:[Node:root|Parent:None]
Child is self
Child not added. Recursive loop.
Tried to add: [[Node:root|Parent:None]] to Self:[[Node:child1|Parent:[Node:root|Parent:None]]] | Ancestors:[Node:child1|Parent:[Node:root|Parent:None]]
Child is ancestor

---TEST METHODS---

roots children:  [[Node:child1|Parent:[Node:root|Parent:None]], [Node:child2|Parent:[Node:root|Parent:None]]]
child1 children:  [[Node:child2|Parent:[Node:root|Parent:None]], [Node:child2|Parent:[Node:root|Parent:None]]]
root parent:  None
child1 parent:  [Node:root|Parent:None]
child2 parent:  [Node:root|Parent:None]
root siblings:  None
child1 siblings:  [[Node:child2|Parent:[Node:root|Parent:None]]]
child1 index:  0
child2 index:  1

---TEST REMOVE CHILD---

roots children: [[Node:child1|Parent:[Node:root|Parent:None]], [Node:child2|Parent:[Node:root|Parent:None]]]
remove child: [[Node:child1|Parent:[Node:root|Parent:None]]]
roots children: [[Node:child1|Parent:[Node:root|Parent:None]]]
[Finished in 0.1s]

Answer 1

It's not great that the node is trying to manage the integrity of the tree, with those checks to prevent adding ancestors as children and the other related checks. It would be better if a node was dumber, and a tree class was in charge of the integrity. For the simple reason of separating concerns, the single responsibility principle.

Another big problem with the current integrity checks is that they are incomplete. They don't protect against adding a sibling as a child, which would "break" the tree. As such, the checks may give a false sense of integrity, which is arguably worse than no checks at all and this no misleading either.

When checking if a node is an ancestor, you build a list of all ancestors and then do a is-in-the-list check. This is a poor approach for several reasons:

• the list of all ancestors takes space
• you could stop adding ancestors when you find a matching node
• the final is-in-the-list operation has linear complexity

You could solve all these issues by not building a list, but traversing the tree until you find a match and then short-circuit (stop searching), or reach the root. The implementation of such logic would fit nicely into a tree class.

Answer 2

Python style doesn't favor the use of accessor functions (the get_xxx function). You could eliminate several of those.

The function get_ancestors cries out for an iterator, something like:

def iter_ancestors(self):
    a = self
    while a.parent is not None:
        a = a.parent
        yield a

If you want a list of ancestors then you do this:

def ancestors(self):
    return list(self.iter_ancestors())

Your function get_sibblings returns None if the Node lacks a parent but returns a list if it does. It should return an empty list in the first case. Also the function name is mis-spelled.

I would suggest this instead:

def siblings(self):
    if self.parent is None:
        return []
    else:
        return [a for a in self.parent.children if a is not self]

I also note that you used != None in the get_ancestors function. That should be is not None instead.