Improve leaves and pruning

This commit is contained in:
Ethan Dalool 2017-04-01 00:10:38 -07:00
parent 97242dab2d
commit f9ccb03a07

View file

@ -42,6 +42,9 @@ OPERATOR_FUNCTIONS = {
'NOT': func_not, 'NOT': func_not,
} }
class NoTokens(Exception):
pass
class ExpressionTree: class ExpressionTree:
def __init__(self, token, parent=None): def __init__(self, token, parent=None):
self.children = [] self.children = []
@ -87,6 +90,9 @@ class ExpressionTree:
if isinstance(tokens, str): if isinstance(tokens, str):
tokens = tokenize(tokens) tokens = tokenize(tokens)
if tokens == []:
raise NoTokens()
if isinstance(tokens[0], list): if isinstance(tokens[0], list):
current = cls.parse(tokens[0], spaces=spaces+1) current = cls.parse(tokens[0], spaces=spaces+1)
else: else:
@ -194,20 +200,31 @@ class ExpressionTree:
raise override from e raise override from e
raise e raise e
@property
def is_leaf(self):
return self.token not in OPERATORS
def map(self, function): def map(self, function):
for node in self.walk(): '''
if node.token in OPERATORS: Apply this function to all of the operands.
continue '''
for node in self.walk_leaves():
node.token = function(node.token) node.token = function(node.token)
def prune(self): def prune(self):
''' '''
Remove any nodes where `token` is None. Remove any nodes where `token` is None.
''' '''
self.children = [child for child in self.children if child is not None] self.children = [child for child in self.children if child.token is not None]
for child in self.children: for child in self.children:
child.prune() child.prune()
if self.token in OPERATORS and len(self.children) == 0:
self.token = None
if self.parent is not None:
self.parent.children.remove(self)
def rootmost(self): def rootmost(self):
current = self current = self
while current.parent is not None: while current.parent is not None:
@ -221,7 +238,7 @@ class ExpressionTree:
def walk_leaves(self): def walk_leaves(self):
for node in self.walk(): for node in self.walk():
if node.token not in OPERATORS: if node.is_leaf:
yield node yield node