split movie guesser into video+movie guesser
[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 logging
23
24 log = logging.getLogger("guessit.guess")
25
26
27 class Guess(dict):
28     def __init__(self, *args, **kwargs):
29         try:
30             confidence = kwargs.pop('confidence')
31         except KeyError:
32             confidence = 0
33
34         dict.__init__(self, *args, **kwargs)
35
36         self._confidence = {}
37         for prop in self:
38             self._confidence[prop] = confidence
39
40     def to_json(self):
41         parts = json.dumps(self, indent = 4).split('\n')
42         for i, p in enumerate(parts):
43             if p[:5] != '    "':
44                 continue
45
46             prop = p.split('"')[1]
47             parts[i] = ('    [%.2f] ' % self._confidence.get(prop, -1)) + p[5:]
48
49         return '\n'.join(parts)
50
51     def confidence(self, prop):
52         return self._confidence[prop]
53
54     def set(self, prop, value, confidence = None):
55         self[prop] = value
56         if confidence is not None:
57             self._confidence[prop] = confidence
58
59     def set_confidence(self, prop, value):
60         self._confidence[prop] = value
61
62     def update(self, other, confidence = None):
63         dict.update(self, other)
64         if isinstance(other, Guess):
65             for prop in other:
66                 self._confidence[prop] = other.confidence(prop)
67
68         if confidence is not None:
69             for prop in other:
70                 self._confidence[prop] = confidence
71
72     def update_highest_confidence(self, other):
73         """Update this guess with the values from the given one. In case there is
74         property present in both, only the one with the highest one is kept."""
75         if not isinstance(other, Guess):
76             raise ValueError, 'Can only call this function on Guess instances'
77
78         for prop in other:
79             if prop in self and self._confidence[prop] >= other._confidence[prop]:
80                 continue
81             self[prop] = other[prop]
82             self._confidence[prop] = other._confidence[prop]
83
84
85
86
87 def choose_int(g1, g2):
88     """Function used by merge_similar_guesses to choose between 2 possible properties
89     when they are integers."""
90     v1, c1 = g1 # value, confidence
91     v2, c2 = g2
92     if (v1 == v2):
93         return (v1, 1 - (1-c1)*(1-c2))
94     else:
95         if c1 > c2:
96             return (v1, c1 - c2)
97         else:
98             return (v2, c2 - c1)
99
100 def choose_string(g1, g2):
101     """Function used by merge_similar_guesses to choose between 2 possible properties
102     when they are strings."""
103     v1, c1 = g1 # value, confidence
104     v2, c2 = g2
105     v1, v2 = v1.strip(), v2.strip()
106     v1l, v2l = v1.lower(), v2.lower()
107
108     combined_prob = 1 - (1-c1)*(1-c2)
109
110     if v1l == v2l:
111         return (v1, combined_prob)
112
113     # check for common patterns
114     elif v1l == 'the ' + v2l:
115         return (v1, combined_prob)
116     elif v2l == 'the ' + v1l:
117         return (v2, combined_prob)
118
119     # if one string is contained in the other, return the shortest one
120     elif v2l in v1l:
121         return (v2, combined_prob)
122     elif v1l in v2l:
123         return (v1, combined_prob)
124
125     # in case of conflict, return the one with highest priority
126     else:
127         if c1 > c2:
128             return (v1, c1 - c2)
129         else:
130             return (v2, c2 - c1)
131
132
133 def merge_similar_guesses(guesses, prop, choose):
134     """Take a list of guesses and merge those which have the same properties,
135     increasing or decreasing the confidence depending on whether their values
136     are similar."""
137
138     similar = [ guess for guess in guesses if prop in guess ]
139     if len(similar) < 2:
140         # nothing to merge
141         return
142
143     if len(similar) > 2:
144         log.warning('merge too complex to be dealt with at the moment, bailing out...')
145         return
146
147     g1, g2 = similar
148
149     if len(set(g1) & set(g2)) > 1:
150         log.warning('both guesses to be merged have more than one property in common, bailing out...')
151         return
152
153     # merge all props of s2 into s1, updating the confidence for the considered property
154     v1, v2 = g1[prop], g2[prop]
155     c1, c2 = g1.confidence(prop), g2.confidence(prop)
156
157     new_value, new_confidence = choose((v1, c1), (v2, c2))
158     if new_confidence >= c1:
159         log.debug("Updating matching property '%s' with confidence %.2f" % (prop, new_confidence))
160     else:
161         log.debug("Updating non-matching property '%s' with confidence %.2f" % (prop, new_confidence))
162
163     g2[prop] = new_value
164     g2.set_confidence(prop, new_confidence)
165
166     g1.update(g2)
167     guesses.remove(g2)
168
169
170 def merge_all(guesses):
171     """Merges all the guesses in a single result and returns it."""
172     result = guesses[0]
173
174     for g in guesses[1:]:
175         if set(result) & set(g):
176             log.warning('duplicate properties %s in merged result...' % (set(result) & set(g)))
177         result.update_highest_confidence(g)
178
179     return result
180