Nonbinary Tree Traversal Method

Question

I am trying to create a python nonbinary tree class with methods that will allow me to get specific nodes, and apply other methods to those found nodes.

I have started with code from this very nice example: https://www.youtube.com/watch?v=4r_XR9fUPhQ

I have added a method that accepts a string that corresponds to the node ("/" separated, skipping the root) I am looking for, recursively searches the tree, and in theory, returns the node using "self" so I can apply another method on it.

However, when I return(self) it gives me a nonetype back instead of the node.

Advice on how to fix this, or a suggestion of another way to go if this is a bad way to structure things would be greatly appreciated!

Thanks in Advance.

Note: this is only set up to match on leaves so far, but I can fix that if I can get the darn thing to return the node I want.

Code Below:

class TreeNode:
    def __init__(self, data):
        self.data = data
        self.children = []
        self.parent = None
        self.forecast = None

    def get_level(self):
        level = 0
        p = self.parent
        while p:
            level += 1
            p = p.parent

        return level

    def print_tree(self):
        spaces = ' ' * self.get_level() * 3
        prefix = spaces + "|__" if self.parent else ""
        print(prefix + self.data)
        if self.children:
            for child in self.children:
                child.print_tree()


    def get_node(self, path):
        segs = path.split('/')
        sep = "/"
        print(self.return_child_names())
        if self.children:
            for child in self.return_child_names():
                if segs[0] == child:
                    found_child = segs.pop(0)
                    break
            self.children[self.return_child_names().index(found_child)].get_node(sep.join(segs))
        else:
            print("Found the node!")
            print(self)
            print(self.data)
            return(self)

    def return_child_names(self):
        return([c.data for c in self.children])

    def add_child(self, child):
        child.parent = self
        self.children.append(child)




def build_product_tree():
    root = TreeNode("Electronics")

    laptop = TreeNode("Laptop")
    laptop.add_child(TreeNode("Mac"))
    laptop.add_child(TreeNode("Surface"))
    laptop.add_child(TreeNode("Thinkpad"))

    cellphone = TreeNode("Cell Phone")
    cellphone.add_child(TreeNode("iPhone"))
    cellphone.add_child(TreeNode("Google Pixel"))
    cellphone.add_child(TreeNode("Vivo"))

    tv = TreeNode("TV")
    tv.add_child(TreeNode("Samsung"))
    tv.add_child(TreeNode("LG"))

    root.add_child(laptop)
    root.add_child(cellphone)
    root.add_child(tv)

    root.print_tree()

    return(root)

product_tree = build_product_tree()

product_tree.get_node("Laptop/Mac").print_tree()

Answer 1

The problem is in the function get_node: you are not always returning a value. Notably in the following block there is no return, and therefore the function can return None:

if self.children:
    for child in self.return_child_names():
        if segs[0] == child:
            found_child = segs.pop(0)
            break
    self.children[self.return_child_names().index(found_child)].get_node(sep.join(segs))

The value you get from the recursive call to get_node is neglected. You should actually return it:

    return self.children[self.return_child_names().index(found_child)].get_node(sep.join(segs))