a few more fixes here and there
[smewt:guessit.git] / guessit / guess.py
1 #!/usr/bin/env python
2 # -*- coding: utf-8 -*-
3 #
4 # GuessIt - A library for guessing information from filenames
5 # Copyright (c) 2011 Nicolas Wack <wackou@gmail.com>
6 #
7 # GuessIt is free software; you can redistribute it and/or modify it under
8 # the terms of the Lesser GNU General Public License as published by
9 # the Free Software Foundation; either version 3 of the License, or
10 # (at your option) any later version.
11 #
12 # GuessIt is distributed in the hope that it will be useful,
13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15 # Lesser GNU General Public License for more details.
16 #
17 # You should have received a copy of the Lesser GNU General Public License
18 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
19 #
20
21 import json
22 import datetime
23 import logging
24
25 log = logging.getLogger("guessit.guess")
26
27
28 class Guess(dict):
29     """A Guess is a dictionary which has an associated confidence for each of its values.
30
31     As it is a subclass of dict, you can use it everywhere you expect a simple dict"""
32     def __init__(self, *args, **kwargs):
33         try:
34             confidence = kwargs.pop('confidence')
35         except KeyError:
36             confidence = 0
37
38         dict.__init__(self, *args, **kwargs)
39
40         self._confidence = {}
41         for prop in self:
42             self._confidence[prop] = confidence
43
44     def to_utf8_dict(self):
45         from guessit.language import Language
46         data = dict(self)
47         for prop, value in data.items():
48             if isinstance(value, datetime.date):
49                 data[prop] = value.isoformat()
50             elif isinstance(value, Language):
51                 data[prop] = str(value)
52             elif isinstance(value, unicode):
53                 data[prop] = value.encode('utf-8')
54             elif isinstance(value, list):
55                 data[prop] = [ str(x) for x in value ]
56
57         return data
58
59     def to_json(self):
60         """NB: this doesn't return a valid json, maybe it should be renamed..."""
61         data = self.to_utf8_dict()
62
63         parts = json.dumps(data, indent = 4).split('\n')
64         for i, p in enumerate(parts):
65             if p[:5] != '    "':
66                 continue
67
68             prop = p.split('"')[1]
69             parts[i] = ('    [%.2f] "' % (self._confidence.get(prop) or -1)) + p[5:]
70
71         return '\n'.join(parts)
72
73     def __str__(self):
74         return str(self.to_utf8_dict())
75
76     def confidence(self, prop):
77         return self._confidence[prop]
78
79     def set(self, prop, value, confidence = None):
80         self[prop] = value
81         if confidence is not None:
82             self._confidence[prop] = confidence
83
84     def set_confidence(self, prop, value):
85         self._confidence[prop] = value
86
87     def update(self, other, confidence = None):
88         dict.update(self, other)
89         if isinstance(other, Guess):
90             for prop in other:
91                 self._confidence[prop] = other.confidence(prop)
92
93         if confidence is not None:
94             for prop in other:
95                 self._confidence[prop] = confidence
96
97     def update_highest_confidence(self, other):
98         """Update this guess with the values from the given one. In case there is
99         property present in both, only the one with the highest one is kept."""
100         if not isinstance(other, Guess):
101             raise ValueError, 'Can only call this function on Guess instances'
102
103         for prop in other:
104             if prop in self and self._confidence[prop] >= other._confidence[prop]:
105                 continue
106             self[prop] = other[prop]
107             self._confidence[prop] = other._confidence[prop]
108
109
110
111
112 def choose_int(g1, g2):
113     """Function used by merge_similar_guesses to choose between 2 possible properties
114     when they are integers."""
115     v1, c1 = g1 # value, confidence
116     v2, c2 = g2
117     if (v1 == v2):
118         return (v1, 1 - (1-c1)*(1-c2))
119     else:
120         if c1 > c2:
121             return (v1, c1 - c2)
122         else:
123             return (v2, c2 - c1)
124
125 def choose_string(g1, g2):
126     """Function used by merge_similar_guesses to choose between 2 possible properties
127     when they are strings."""
128     v1, c1 = g1 # value, confidence
129     v2, c2 = g2
130     v1, v2 = v1.strip(), v2.strip()
131     v1l, v2l = v1.lower(), v2.lower()
132
133     combined_prob = 1 - (1-c1)*(1-c2)
134
135     if v1l == v2l:
136         return (v1, combined_prob)
137
138     # check for common patterns
139     elif v1l == 'the ' + v2l:
140         return (v1, combined_prob)
141     elif v2l == 'the ' + v1l:
142         return (v2, combined_prob)
143
144     # if one string is contained in the other, return the shortest one
145     elif v2l in v1l:
146         return (v2, combined_prob)
147     elif v1l in v2l:
148         return (v1, combined_prob)
149
150     # in case of conflict, return the one with highest priority
151     else:
152         if c1 > c2:
153             return (v1, c1 - c2)
154         else:
155             return (v2, c2 - c1)
156
157
158 def _merge_similar_guesses_nocheck(guesses, prop, choose):
159     """Take a list of guesses and merge those which have the same properties,
160     increasing or decreasing the confidence depending on whether their values
161     are similar.
162
163     This function assumes there are at least 2 valid guesses."""
164
165     similar = [ guess for guess in guesses if prop in guess ]
166
167     g1, g2 = similar[0], similar[1]
168
169     if len(set(g1) & set(g2)) > 1:
170         log.warning('both guesses to be merged have more than one property in common, bailing out...')
171         return
172
173     # merge all props of s2 into s1, updating the confidence for the considered property
174     v1, v2 = g1[prop], g2[prop]
175     c1, c2 = g1.confidence(prop), g2.confidence(prop)
176
177     new_value, new_confidence = choose((v1, c1), (v2, c2))
178     if new_confidence >= c1:
179         log.debug("Updating matching property '%s' with confidence %.2f" % (prop, new_confidence))
180     else:
181         log.debug("Updating non-matching property '%s' with confidence %.2f" % (prop, new_confidence))
182
183     g2[prop] = new_value
184     g2.set_confidence(prop, new_confidence)
185
186     g1.update(g2)
187     guesses.remove(g2)
188
189 def merge_similar_guesses(guesses, prop, choose):
190     """Take a list of guesses and merge those which have the same properties,
191     increasing or decreasing the confidence depending on whether their values
192     are similar."""
193
194     similar = [ guess for guess in guesses if prop in guess ]
195     if len(similar) < 2:
196         # nothing to merge
197         return
198
199     if len(similar) == 2:
200         _merge_similar_guesses_nocheck(guesses, prop, choose)
201
202     if len(similar) > 2:
203         log.debug('complex merge, trying our best...')
204         _merge_similar_guesses_nocheck(guesses, prop, choose)
205         merge_similar_guesses(guesses, prop, choose)
206         return
207
208
209 def merge_append_guesses(guesses, prop):
210     """Take a list of guesses and merge those which have the same properties by
211     appending them in a list."""
212
213     similar = [ guess for guess in guesses if prop in guess ]
214     if not similar:
215         return
216
217     merged = similar[0]
218     merged[prop] = [ merged[prop] ]
219     # TODO: what to do with global confidence? mean of them all?
220
221     for m in similar[1:]:
222         for prop2 in m:
223             if prop == prop2:
224                 merged[prop].append(m[prop])
225             else:
226                 if prop2 in m:
227                     log.warning('overwriting property "%s" with value ' % (prop2, m[prop2]))
228                 merged[prop2] = m[prop2]
229                 # TODO: confidence also
230
231         guesses.remove(m)
232
233
234 def merge_all(guesses, append = []):
235     """Merges all the guesses in a single result, removes very unlikely values, and returns it.
236     You can specify a list of properties that should be appended into a list instead of being
237     merged.
238
239     >>> merge_all([ Guess({ 'season': 2 }, confidence = 0.6),
240     ...             Guess({ 'episodeNumber': 13 }, confidence = 0.8) ])
241     {'season': 2, 'episodeNumber': 13}
242
243     >>> merge_all([ Guess({ 'episodeNumber': 27 }, confidence = 0.02),
244     ...             Guess({ 'season': 1 }, confidence = 0.2) ])
245     {'season': 1}
246
247     """
248     if not guesses:
249         return Guess()
250
251     result = guesses[0]
252
253     for g in guesses[1:]:
254         # first append our appendable properties
255         for prop in append:
256             if prop in g:
257                 result.set(prop, result.get(prop, []) + [ g[prop] ],
258                            # TODO: what to do with confidence here? maybe an arithmetic mean...
259                            confidence = g.confidence(prop))
260
261                 del g[prop]
262
263         # then merge the remaining ones
264         if set(result) & set(g):
265             log.warning('duplicate properties %s in merged result...' % (set(result) & set(g)))
266
267         result.update_highest_confidence(g)
268
269     # delete very unlikely values
270     for p in result.keys():
271         if result.confidence(p) < 0.05:
272             del result[p]
273
274     # make sure our appendable properties contain unique values
275     for prop in append:
276         if prop in result:
277             result[prop] = list(set(result[prop]))
278
279     return result
280