Extract decision_tree_predict().

darius · darius · commit 305e70299519 · 2011-10-23T18:55:41.000Z
diff --git a/learning.py b/learning.py
@@ -257,11 +257,8 @@ def __init__(self, attr, attrname=None, branches=None):
 
     def predict(self, example):
         "Given an example, use the tree to classify the example."
-        child = self.branches[example[self.attr]]
-        if isinstance(child, DecisionTree):
-            return child.predict(example)
-        else:
-            return child
+        attrvalue = example[self.attr]
+        return decision_tree_predict(self.branches[attrvalue], example)
 
     def add(self, val, subtree):
         "Add a branch.  If self.attr = val, go to the given subtree."
@@ -280,19 +277,18 @@ def display(self, indent=0):
     def __repr__(self):
         return ('DecisionTree(%r, %r, %r)'
                 % (self.attr, self.attrname, self.branches))
-
-Yes, No = True, False
+    
+def decision_tree_predict(tree, example):
+    "Treat a non-DecisionTree as a leaf."
+    return tree.predict(example) if isinstance(tree, DecisionTree) else tree
 
 #______________________________________________________________________________
 
 class DecisionTreeLearner(Learner):
     "[Fig. 18.5]"
 
     def predict(self, example):
-        if isinstance(self.dt, DecisionTree):
-            return self.dt.predict(example)
-        else:
-            return self.dt
+        return decision_tree_predict(self.dt, example)
 
     def train(self, dataset):
         self.dataset = dataset
