Add epa data to all tests where it works

Author: dankegel

libvin/epa.py

@@ -25,8 +25,9 @@ def __init__(self, vin):
         self.__nhtsa = nhtsa_decode(vin)
         self.__attribs = self.__get_attributes()
         self.__model = self.__get_model()
-        (self.__ids, self.__trims) = self.__get_ids()
-        self.__eco   = self.__get_vehicle_economy()
+        if (self.__model != None):
+            self.__ids, self.__trims = self.__get_ids()
+            self.__eco   = self.__get_vehicle_economy()
 
     @property
     def nhtsa(self):
@@ -58,6 +59,15 @@ def id(self):
         # first one.  We're guessing anyway, what with the fuzzy matching and all...
         return self.__ids[0]
 
+    @property
+    def trim(self):
+        '''
+        EPA trim for this vehicle.
+        '''
+        # FIXME: If we don't know which trim exactly, just pick the
+        # first one.  We're guessing anyway, what with the fuzzy matching and all...
+        return self.__trims[0]
+
     @property
     def eco(self):
         '''
@@ -80,7 +90,7 @@ def __get_attributes(self):
         driveType = self.nhtsa['DriveType']
         if 'AWD' in driveType:
             attributes.append("AWD")
-        if '4WD' in driveType or '4x4' in driveType:
+        elif '4WD' in driveType or '4x4' in driveType:
             attributes.append("4WD")
         elif '4x2' in driveType:
             attributes.append("2WD")
@@ -94,9 +104,16 @@ def __get_attributes(self):
             attributes.append(self.nhtsa['BodyClass'])
         if 'Series' in self.nhtsa and self.nhtsa['Series'] != "":
             attributes.append(self.nhtsa['Series'])
+        if 'Series2' in self.nhtsa and self.nhtsa['Series2'] != "":
+            attributes.append(self.nhtsa['Series2'])
 
         if 'DisplacementL' in self.nhtsa and self.nhtsa['DisplacementL'] != '':
             attributes.append('%s L' % self.nhtsa['DisplacementL'])
+            # EPA sometimes likes to go all precise
+            if '.' not in self.nhtsa['DisplacementL']:
+               attributes.append('%s.0 L' % self.nhtsa['DisplacementL'])
+        if 'EngineCylinders' in self.nhtsa and self.nhtsa['EngineCylinders'] != '':
+            attributes.append('%s cyl' % self.nhtsa['EngineCylinders'])
 
         if 'Manual' in self.nhtsa['TransmissionStyle']:
             attributes.append('MAN')
@@ -106,6 +123,10 @@ def __get_attributes(self):
             attributes.append('CVT')
             attributes.append('Variable')
 
+        # Twin turbo is "Yes, Yes"!
+        if 'Turbo' in self.nhtsa and 'Yes' in self.nhtsa['Turbo']:
+            attributes.append('Turbo')
+
         return attributes
 
     def __get_possible_models(self):
@@ -171,40 +192,45 @@ def __get_possible_ids(self):
             return None
         return id2trim
 
-    def __fuzzy_match(self, attributes, choices):
+    def __fuzzy_match(self, mustmatch, attributes, choices):
         '''
         Given a base name and a bunch of attributes, find the choice that matches them the best.
-        name : string
+        mustmatch : string
         attributes : string[]
         choices : dict mapping id to string
         Returns: array of ids of best matching choices
         '''
 
-        best_matches = 0
-        best_ids = []
-        best_fraction = 0
+        best_ids = []     # id of best matching trims
+        best_len = 0      # len of best matching trims
+        best_matched = 0
         for (key, val) in choices.iteritems():
-            n_matches = 0
-            chars_total = len(val)
+            # optional mandatory attribute
+            # to prevent [Q60 AWD] from matching Q85 AWD instead of Q60 AWD Coupe
+            if mustmatch != None and mustmatch.upper() not in val.upper():
+                continue
+            # Find choice that matches most chars from attributes.
+            # In case of a tie, prefer shortest choice.
             chars_matched = 0
             for attrib in attributes:
                 if attrib != "" and attrib.upper() in val.upper():
-                    chars_matched += len(attrib)
-                    n_matches += 1
-            fraction = float(chars_matched) / chars_total
-            #print "n_matches %d, chars_matched %d, chars_total %d, fraction %f for %s" % (n_matches, chars_matched,chars_total, fraction, val)
-            if (n_matches > best_matches):
+                    if chars_matched == 0:
+                        chars_matched = len(attrib)
+                    else:
+                        chars_matched += len(attrib) + 1  # for space
+            #print "chars_matched %d, for %s" % (chars_matched, val)
+            if (chars_matched > best_matched):
                 best_ids = [key]
-                best_matches = n_matches
-                best_fraction = fraction
-            elif ((n_matches > 0) and (n_matches == best_matches)):
-                # Heuristic: favor most complete match (so missing an attribute like Hybrid hurts)
-                if fraction > best_fraction:
+                best_len = len(val)
+                best_matched = chars_matched
+            elif (chars_matched > 0 and chars_matched == best_matched):
+                if len(val) < best_len:
+                    #print "chars %d == %d, len %d < %d, breaking tie in favor of shorter trim" % (chars_matched, best_matched, len(val), best_len)
                     best_ids = [key]
-                    best_matches = n_matches
-                    best_fraction = fraction
-                elif fraction == best_fraction:
-                    #print "%d == %d, marking tie" % (n_matches, best_matches)
+                    best_len = len(val)
+                    best_matched = chars_matched
+                elif len(val) == best_len:
+                    #print "chars %d == %d, len %d == %d, marking tie" % (chars_matched, best_matched, len(val), best_len)
                     best_ids.append(key)
         if len(best_ids) == 0:
             print "epa:__fuzzy_match: no match found for vin %s" % self.vin
@@ -220,7 +246,31 @@ def __get_model(self):
         id2models = self.__get_possible_models()
         if id2models == None:
             return None
-        ids = self.__fuzzy_match(self.__attribs, id2models)
+        #print "Finding model for vin %s" % self.vin
+        # Special case for Mercedes-Benz, which puts the real model in Series
+        oldmodel = self.nhtsa['Model']
+        model = oldmodel.replace('-Class', '')
+        ids = self.__fuzzy_match(model, self.__attribs, id2models)
+        if len(ids) != 1:
+            # Second chance for alternate spellings
+            if '4WD' in self.__attribs:
+                tribs = self.__attribs
+                tribs.append('AWD')
+                print "Searching again with AWD"
+                ids = self.__fuzzy_match(self.nhtsa['Model'], tribs, id2models)
+            elif '2WD' in self.__attribs and 'FWD' not in self.__attribs:
+                tribs = self.__attribs
+                tribs.append('RWD')
+                print "Searching again with RWD"
+                ids = self.__fuzzy_match(self.nhtsa['Model'], tribs, id2models)
+            elif 'Mazda' in self.nhtsa['Model']:
+                oldmodel = self.nhtsa['Model']
+                model = oldmodel.replace('Mazda', '')
+                tribs = self.__attribs
+                tribs.append(model)
+                print "Searching again with %s instead of %s" % (model, oldmodel)
+                ids = self.__fuzzy_match(model, tribs, id2models)
+
         if len(ids) != 1:
             print "epa:__get_model: Failed to find model for vin %s" % self.vin
             return None
@@ -238,7 +288,8 @@ def __get_ids(self):
         id2trim = self.__get_possible_ids()
         if id2trim == None:
             return None
-        ids = self.__fuzzy_match(self.__attribs, id2trim)
+        #print "Finding trims for vin %s" % self.vin
+        ids = self.__fuzzy_match(None, self.__attribs, id2trim)
         if len(ids) == 0:
             print "epa:__get_id: No trims found for vin %s" % self.vin
             return None