Source Code for Module teamwork.test.math.testID3

 1  from teamwork.math.dtree import * 
 2  from teamwork.math.id3 import * 
 3  import sys 
 4   
 5   
 6 -def get_file(): 
 7      """ 
 8      Tries to extract a filename from the command line.  If none is present, it 
 9      prompts the user for a filename and tries to open the file.  If the file  
10      exists, it returns it, otherwise it prints an error message and ends 
11      execution.  
12      """ 
13      # Get the name of the data file and load it into  
14      if len(sys.argv) < 2: 
15          # Ask the user for the name of the file 
16          print "Filename: ",  
17          filename = sys.stdin.readline().strip() 
18      else: 
19          filename = sys.argv[1] 
20   
21      try: 
22          fin = open(filename, "r") 
23      except IOError: 
24          print "Error: The file '%s' was not found on this system." % filename 
25          sys.exit(0) 
26   
27      return fin 
28   
29 -def run_test(fin): 
30      """ 
31      This function creates a list of exmaples data (used to learn the d-tree) 
32      and a list of samples (for classification by the d-tree) from the 
33      designated file.  It then creates the d-tree and uses it to classify the 
34      samples.  It prints the classification of each record in the samples list 
35      and returns the d-tree. 
36      """ 
37      # Create a list of all the lines in the data file 
38      lines = [line.strip() for line in fin.readlines()] 
39   
40      # Remove the attributes from the list of lines and create a list of 
41      # the attributes. 
42      lines.reverse() 
43      attributes = [attr.strip() for attr in lines.pop().split(",")] 
44      target_attr = attributes[-1] 
45      lines.reverse() 
46   
47      # Create a list of the data in the data file 
48      data = [] 
49      for line in lines: 
50          data.append(dict(zip(attributes, 
51                               [datum.strip() for datum in line.split(",")]))) 
52           
53      # Copy the data list into the examples list for testing 
54      examples = data[:] 
55      # Create the decision tree 
56      tree = create_decision_tree(data, attributes, target_attr, gain) 
57   
58      # Classify the records in the examples list 
59      classification = classify(tree, examples) 
60   
61      # Print out the classification for each record 
62      for item in classification: 
63          print item 
64   
65      return tree 
66   
67 -def print_tree(tree, prefix= ''): 
68      """ 
69      This function recursively crawls through the d-tree and prints it out in a 
70      more readable format than a straight print of the Python dict object.   
71      """ 
72      if type(tree) == dict: 
73          print "%s%s" % (prefix, tree.keys()[0]) 
74          for item in tree.values()[0].keys(): 
75              print "%s\t%s" % (prefix, item) 
76              print_tree(tree.values()[0][item], prefix + "\t") 
77      else: 
78          print "%s\t->\t%s" % (prefix, tree) 
79   
80   
81  if __name__ == "__main__": 
82      fin = get_file() 
83      print "------------------------\n" 
84      print "--   Classification   --\n" 
85      print "------------------------\n" 
86      print "\n"     
87      tree = run_test(fin) 
88      print "\n" 
89      print "------------------------\n" 
90      print "--   Decision Tree    --\n" 
91      print "------------------------\n" 
92      print "\n" 
93      print_tree(tree, "") 
94      fin.close() 
95