class ExistingChild(Exception): pass class InvalidIdentifier(Exception): pass class Tree: def __init__(self, identifier, data=None): self.assert_identifier_ok(identifier) self.identifier = identifier self.data = data self.parent = None self.children = {} def __eq__(self, other): return isinstance(other, Tree) and self.abspath() == other.abspath() def __getitem__(self, key): return self.children[key] def __hash__(self): return hash(self.abspath()) def __repr__(self): return f'Tree({self.identifier})' @staticmethod def assert_identifier_ok(identifier): if not isinstance(identifier, str): raise InvalidIdentifier(f'Identifier {identifier} must be a string.') if '/' in identifier or '\\' in identifier: raise InvalidIdentifier('Identifier cannot contain slashes') def abspath(self): node = self nodes = [node] while nodes[-1].parent is not None: nodes.append(nodes[-1].parent) nodes.reverse() nodes = [node.identifier for node in nodes] return '\\'.join(nodes) def add_child(self, other_node, overwrite_parent=False): self.assert_child_available(other_node.identifier) if other_node.parent is not None and not overwrite_parent: raise ValueError('That node already has a parent. Try `overwrite_parent=True`') other_node.parent = self self.children[other_node.identifier] = other_node return other_node def assert_child_available(self, identifier): if identifier in self.children: raise ExistingChild(f'Node {self.identifier} already has child {identifier}') def detach(self): if self.parent is None: return del self.parent.children[self.identifier] self.parent = None def list_children(self, sort=None): children = list(self.children.values()) if sort is None: children.sort(key=lambda node: (node.identifier.lower(), node.identifier)) else: children.sort(key=sort) return children def walk(self, sort=None): yield self for child in self.list_children(sort=sort): yield from child.walk(sort=sort) def walk_parents(self): parent = self.parent while parent is not None: yield parent parent = parent.parent