Game Graph with python visitor pattern

main.py

class Player:
    def visit(self, tree_node):
        # tree_node should be of type TreeNode
        print(tree_node.story_piece)
        if not tree_node.choices:
            # returning 'None' will be used to signal the end of traversal
            return None

        # use a lambda to keep it DRY
        prompter = lambda : input("Enter 1 or 2 to continue the story: ")
        choice = prompter()
        while choice not in ["1", "2"]:
            print("Invalid choice. Try again.")
            choice = prompter()
        
        chosen_index = int(choice) - 1
        return tree_node.choices[chosen_index]

def main():
    # initialize the player (as a "Visitor")
    player = Visitor()

    # play
    current_node = story_data
    while current_node:
        current_node = player.visit(current_node)

    # this will make it clear when the game ends as expected
    print("no more choices, game over")

def main():
    pass

if __name__ == '__main__':
    main()

story_data.py

from tree import TreeNode

story_data = TreeNode("start- choose:\n1: A,\n2: B")
story_a = TreeNode("A- choose:\n1: A1,\n2: A2")
story_a1 = TreeNode("A1: true/best ending! yay!")
story_a2 = TreeNode("A2: worst ending T.T")
story_b = TreeNode("B: bad ending :(")

# story_a and story_b are outcomes from the root story
story_data.add_child(story_a)
story_data.add_child(story_b)

# story_a1 and story_a2 are outcomes from story_a
story_a.add_child(story_a1)
story_a.add_child(story_a2)

tree.py

class TreeNode:
    def __init__(self, story_piece, choices=None):
        self.story_piece = story_piece
        self.choices = choices or []

    def add_child(self, node):
        self.choices.append(node)

    def play(self, player):
        # allows us to separate gameplay implementation from story implementation
        return player.visit(self)

For demonstrative purposes answering stackoverflow post https://stackoverflow.com/questions/77076704/adventure-story-graph-tuple-bug