Use list.sort's keyword parameters
[joel/kofoto.git] / src / packages / kofoto / gkofoto / cachingpixbufloader.py
1 """This module contains the CachingPixbufLoader class."""
2
3 __all__ = ["CachingPixbufLoader"]
4
5 import gc
6 import os
7 import time
8 if __name__ == "__main__":
9     import pygtk
10     pygtk.require("2.0")
11 import gobject
12 from kofoto.common import UnimplementedError
13 from kofoto.gkofoto.pixbufloader import PixbufLoader, get_pixbuf_size
14 from kofoto.gkofoto.pseudothread import PseudoThread
15 from kofoto.iodict import InsertionOrderedDict
16 from kofoto.rectangle import Rectangle
17
18 class _RequestStateBase(object):
19     def __init__(self, request):
20         self._request = request
21
22     def add_callback(self, load_callback, error_callback):
23         self._request._callbacks.append((load_callback, error_callback))
24
25     def get_number_of_loaded_pixels(self):
26         if self._request._loaded_bytes == 0:
27             return 0
28         else:
29             return (
30                 float(self._request._size_to_load[0]) *
31                 self._request._size_to_load[1] *
32                 self._request._loaded_bytes /
33                 self._request._available_bytes)
34
35     def is_finished(self):
36         raise UnimplementedError
37
38     def load_some_more(self):
39         return 0
40
41 class _RequestStateInitial(_RequestStateBase):
42     def get_number_of_loaded_pixels(self):
43         return 0
44
45     def is_finished(self):
46         return False
47
48     def load_some_more(self):
49         req = self._request
50         try:
51             stat = os.stat(req._path)
52             # TODO: Look for image in diskcache here and load the
53             # cached image instead (if cached image mtime <= original
54             # image mtime?).
55         except OSError:
56             req._state = _RequestStateError(req)
57         else:
58             req._mtime = stat.st_mtime
59             req._available_bytes = stat.st_size
60             req._state = _RequestStateWaitingForSize(req)
61             req._pb_loader.prepare(req._path, req._size_limit)
62         return 0
63
64 class _RequestStateWaitingForSize(_RequestStateBase):
65     def get_number_of_loaded_pixels(self):
66         return 0
67
68     def is_finished(self):
69         return False
70
71     def load_some_more(self):
72         req = self._request
73         req._unreported_bytes += req._pb_loader.load_some_more()
74         original_size = req._pb_loader.get_original_size()
75         if original_size is not None:
76             req._state = _RequestStateLoading(req)
77             req._original_size = original_size
78             if req._size_limit is None:
79                 req._size_to_load = original_size
80             else:
81                 req._size_to_load = Rectangle(*original_size).downscaled_to(
82                     req._size_limit)
83         return 0
84
85 class _RequestStateLoading(_RequestStateBase):
86     def is_finished(self):
87         return False
88
89     def load_some_more(self):
90         req = self._request
91         loaded_bytes = req._pb_loader.load_some_more()
92         if req._unreported_bytes > 0:
93             loaded_bytes += req._unreported_bytes
94             req._unreported_bytes = 0
95         req._loaded_bytes += loaded_bytes
96         if loaded_bytes == 0:
97             # Here it would be possible to check mtime again and
98             # reload the image if mtime has changed. We decided to
99             # ignore this right now because of implementation
100             # difficulties (for example, a negative number of pixels
101             # (watch out so it isn't 0!) must be reported back to
102             # req._cpb_loader in some way).
103
104             req._pixbuf = req._pb_loader.get_pixbuf()
105             if req._pixbuf is None:
106                 req._state = _RequestStateError(req)
107             else:
108                 req._state = _RequestStateFinished(req)
109         return (
110             (float(loaded_bytes) / req._available_bytes) *
111             req._size_to_load[0] * req._size_to_load[1])
112
113 class _RequestStateFinished(_RequestStateBase):
114     def __init__(self, request):
115         _RequestStateBase.__init__(self, request)
116         for (load_cb, _) in request._callbacks:
117             if request._cpb_loader._debug_level > 0:
118                 print "%.3f callback(%s, %s)" % (
119                     time.time(), request._path, request._size_limit)
120             load_cb(request._pixbuf, request._original_size)
121         request._callbacks = []
122         # TODO: Store in disk cache here if
123         # self._request._persistence_size_limit isn't None and we
124         # didn't load the image from the cache.
125
126     def add_callback(self, load_callback, _):
127         req = self._request
128         if req._cpb_loader._debug_level > 0:
129             print "%.3f callback(%s, %s)" % (
130                 time.time(), req._path, req._size_limit)
131         load_callback(req._pixbuf, req._original_size)
132
133     def is_finished(self):
134         return True
135
136 class _RequestStateError(_RequestStateBase):
137     def __init__(self, request):
138         _RequestStateBase.__init__(self, request)
139         request._pb_loader.cancel()
140         request._pixbuf = None
141         for (_, error_cb) in request._callbacks:
142             if error_cb is not None:
143                 error_cb()
144         request._callbacks = []
145
146         # Don't cache negative requests since we can't be sure that it
147         # will fail in the future. The image may for example be partly
148         # written now but fully written soon.
149         key = (request._path, request._size_limit)
150         request._cpb_loader._remove_erroneous_request(key)
151
152     def add_callback(self, _, error_callback):
153         error_callback()
154
155     def is_finished(self):
156         return True
157
158 class _Request(object):
159     def __init__(self, cpb_loader, path, size_limit, persistence_size_limit):
160         self._cpb_loader = cpb_loader
161         self._path = path
162         self._size_limit = size_limit
163         self._persistence_size_limit = persistence_size_limit
164         self._callbacks = [] # List of (load_callback, error_callback).
165         self._pixbuf = None
166         self._unreported_bytes = 0
167         self._loaded_bytes = 0
168         self._available_bytes = None
169         self._original_size = None
170         self._size_to_load = None
171         self._mtime = None
172         self._pb_loader = PixbufLoader()
173         self._state = _RequestStateInitial(self)
174
175     # ----------------------------------
176
177     def get_key(self):
178         return (self._path, self._size_limit)
179     key = property(get_key)
180
181     # ----------------------------------
182
183     def add_callback(self, load_callback, error_callback):
184         return self._state.add_callback(load_callback, error_callback)
185
186     def cancel(self):
187         self._pb_loader.cancel()
188         del self._state # Break cycle.
189
190     def get_number_of_loaded_pixels(self):
191         return self._state.get_number_of_loaded_pixels()
192
193     def has_changed_on_disk(self):
194         if self._mtime is None:
195             # We don't know yet.
196             return False
197         else:
198             try:
199                 return os.path.getmtime(self._path) != self._mtime
200             except OSError:
201                 return True
202
203     def is_finished(self):
204         return self._state.is_finished()
205
206     def load_some_more(self):
207         return self._state.load_some_more()
208
209     def remove_callback(self, load_callback, error_callback):
210         try:
211             self._callbacks.remove((load_callback, error_callback))
212         except ValueError:
213             pass
214
215     def _print_state(self, verbose):
216         print "        state:                 ", self._state.__class__.__name__
217         print "        path:                  ", self._path
218         print "        size limit:            ", self._size_limit
219         print "        persistence size limit:", self._persistence_size_limit
220         if verbose:
221             print "        callbacks:             ", self._callbacks
222             print "        pixbuf:                ", self._pixbuf
223             print "        unreported bytes:      ", self._unreported_bytes
224             print "        loaded bytes:          ", self._loaded_bytes
225             print "        available bytes:       ", self._available_bytes
226             print "        original size:         ", self._original_size
227             print "        size to load:          ", self._size_to_load
228             print "        mtime:                 ", self._mtime
229             print "        pb loader:             ", self._pb_loader
230
231 class CachingPixbufLoader(object):
232     """A pixbuf loader with preload and cache functionality.
233
234     This class is a pixbuf loader that keeps loaded pixbufs in an LRU
235     memory cache. It handles several outstanding pixbuf load requests
236     and also handles preload requests. It can optionally store pixbufs
237     in a disk-based cache. Currently, all pixbufs are stored as JPEGs.
238
239     Load requests are asynchronous; the request is done by supplying
240     the path to the image, a wanted pixbuf size and a callback
241     function to the load() method. The callback function is called
242     when the load has finished (or immediately, if the pixbuf already
243     exists in the cache). Old load requests have higher priority than
244     new. Load requests have higher priority than preload requests. 
245     Load requests are never forgotten.
246
247     Preload requests are also asynchronous. A preload request is used
248     to hint the loader to load a pixbuf into the cache. New preload
249     requests have higher priority than old. Load requests have higher
250     priority than preload requests. Old preload requests may be
251     ignored (and the cached pixbuf thrown away) if the cache limits
252     are exceeded.
253     """
254
255     def __init__(self, pixel_limit=10**7, cache_directory=None):
256         """Constructor.
257
258         Arguments:
259
260         pixel_limit  -- The number of pixels to keep in the memory
261                         cache.
262         cache_directory
263                      -- Path to the disk cache directory. If None, no
264                         images will be cached on disk.
265         """
266
267         self._pixel_limit = pixel_limit
268         self._cache_directory = cache_directory
269         self._load_thread = PseudoThread(self._load_loop())
270         self._debug_level = 0
271
272         # Cached value of the sum of loaded pixels of all requests in
273         # the queue.
274         self._pixels_in_cache = 0.0
275
276         # Whether the load loop should reevaluate which request to
277         # work on.
278         self._load_loop_reeval = False
279
280         # Maps path to Rectangle(full_width, full_height).
281         self._available_size = {}
282
283         # Maps (path, (width_limit, height_limit)) to _Request
284         # instances. Newest requests are first and oldest are last.
285         self._request_queue = InsertionOrderedDict()
286
287         # List of _Request instances. Always a subset of
288         # self._request_queue.values() (except order). Newest requests
289         # are last and oldest are first.
290         self._load_queue = []
291
292         # Maps path to set(_Request)
293         self._path_to_requests = {}
294
295     def cancel_load(self, handle):
296         """Cancel a load.
297
298         This method does not remove the cached pixbuf from the cache,
299         it just makes sure that the callback passed to load doesn't
300         get called. In other words, cancel_load converts a load to a
301         preload.
302
303         If the handle represents a load that already has finished,
304         nothing will happen and no exception will be raised.
305         """
306         if self._debug_level > 0:
307             print "%.3f cancel_load(%r)" % (time.time(), handle)
308
309         (path, size_limit, load_callback, error_callback) = handle
310         key = (path, size_limit)
311         if key in self._request_queue:
312             request = self._request_queue[key]
313             try:
314                 self._load_queue.remove(request)
315             except ValueError:
316                 # Nothing to do.
317                 return
318             request.remove_callback(load_callback, error_callback)
319             self._load_loop_reeval = True
320
321     def get_pixel_limit(self):
322         """Get cache size limit."""
323
324         return self._pixel_limit
325
326     def load(self, path, size_limit, load_callback,
327              error_callback=None, persistence_size_limit=None):
328         """Load a pixbuf as quick as possible.
329
330         The pixbuf is taken from the cache if possible.
331
332         Old load requests have higher priority than new. Load requests
333         have higher priority than preload requests.
334
335         Arguments:
336
337         path         -- Path to the image file to load.
338         size_limit   -- A tuple (width, height) with the size limit of
339                         the resulting pixbuf, or None. If None, a
340                         full-size pixbuf will be loaded.
341         load_callback
342                      -- Function to call when loading has finished.
343                         The function will be passed two arguments: the
344                         resulting pixbuf and a tuple (width, height)
345                         with the full size of the image file on disk.
346         error_callback
347                      -- Function to call when loading has failed. If
348                         None, no call will be made and the error will
349                         be ignored. The function will be given no
350                         arguments.
351         persistence_size_limit
352                      -- Limit (an integer) of the image stored in the
353                         disk cache (if disk caching is enabled). Must
354                         be None if size_limit is None. Must be equal
355                         to or greater than max(size_limit[0],
356                         size_limit[1]). If None, the disk cache will
357                         not be used.
358
359         The method returns a handle that can be passed to cancel_load.
360         """
361         if self._debug_level > 0:
362             print "%.3f load(%s, %s)" % (time.time(), path, size_limit)
363
364         if size_limit is not None:
365             size_limit = tuple(size_limit)
366             if persistence_size_limit is not None:
367                 assert persistence_size_limit >= size_limit[0]
368                 assert persistence_size_limit >= size_limit[1]
369             size_limit = self._calculate_size_limit(path, size_limit)
370         key = (path, size_limit)
371         request = self._request_queue.get(key)
372         if request and request.has_changed_on_disk():
373             self._remove_request(key)
374         if key in self._request_queue:
375             # Just move it to the front.
376             request = self._request_queue[key]
377             self._request_queue.insert_first(key, request)
378         else:
379             # Let preload create the request.
380             self.preload(path, size_limit, persistence_size_limit)
381             request = self._request_queue[key]
382         if request not in self._load_queue:
383             self._load_queue.append(request)
384             self._load_loop_reeval = True
385
386         # It's okay to wait until now to add the callback, since any
387         # errors that will trigger error_callback will arise later (at
388         # the first load_some_more in the load loop).
389         request.add_callback(load_callback, error_callback)
390
391         self._load_thread.start()
392         handle = (
393             path, size_limit, load_callback, error_callback)
394         return handle
395
396     def preload(self, path, size_limit, persistence_size_limit=None):
397         """Request that a pixbuf should be loaded into the cache.
398
399         The pixbuf can be retrieved by a call to load().
400
401         New preload requests have higher priority than old. Load
402         requests have higher priority than preload requests.
403
404         Arguments:
405
406         path         -- Path to the image file to load.
407         size_limit   -- A tuple (width, height) with the size limit of the
408                         resulting pixbuf, or None. If None, a
409                         full-size pixbuf will be loaded.
410         persistence_size_limit
411                      -- Limit (an integer) of the image stored in the
412                         disk cache (if disk caching is enabled). Must
413                         be None if size_limit is None. Must be equal
414                         to or greater than max(size_limit[0],
415                         size_limit[1]). If None, the disk cache will
416                         not be used.
417         """
418
419         if self._debug_level > 0:
420             print "%.3f preload(%s, %s)" % (time.time(), path, size_limit)
421
422         if size_limit is not None:
423             size_limit = tuple(size_limit)
424             if persistence_size_limit is not None:
425                 assert persistence_size_limit >= size_limit[0]
426                 assert persistence_size_limit >= size_limit[1]
427             size_limit = self._calculate_size_limit(path, size_limit)
428         key = (path, size_limit)
429         if key in self._request_queue:
430             request = self._request_queue[key]
431         else:
432             request = _Request(self, path, size_limit, persistence_size_limit)
433             request_set = self._path_to_requests.setdefault(path, set())
434             request_set.add(request)
435         self._request_queue.insert_first(key, request)
436         self._load_loop_reeval = True
437         self._load_thread.start()
438
439     def set_pixel_limit(self, pixel_limit):
440         """Set cache size limit.
441
442         Arguments:
443
444         pixel_limit -- The number of pixels to keep in the cache.
445         """
446
447         self._pixel_limit = pixel_limit
448         self._prune_queue()
449
450     def unload(self, path, size_limit):
451         """Remove a cached pixbuf from the cache.
452
453         This method is a hint to the cache that it's not interesting
454         to keep a cached pixbuf anymore. If no pixbuf for the given
455         path and size limit exists, nothing will happen.
456
457         Arguments:
458
459         path         -- Path given to load/preload.
460         size_limit   -- Size limit given to load/preload.
461         """
462
463         if self._debug_level > 0:
464             print "%.3f unload(%s, %s)" % (time.time(), path, size_limit)
465
466         key = (path, size_limit)
467         if key in self._request_queue:
468             request = self._request_queue[key]
469             if request not in self._load_queue:
470                 self._remove_request(key)
471
472     def unload_all(self, path):
473         """Remove all cached sizes of a pixbuf from the cache.
474
475         This method is a hint to the cache that it's not interesting
476         to keep a cached pixbuf anymore. If no pixbufs for the given
477         path exist, nothing will happen.
478
479         Arguments:
480
481         path         -- Path given to load/preload.
482         """
483
484         if self._debug_level > 0:
485             print "%.3f unload_all(%s)" % (time.time(), path)
486
487         if path not in self._path_to_requests:
488             return
489         for request in self._path_to_requests[path].copy():
490             if request not in self._load_queue:
491                 self._remove_request(request.key)
492
493     def _calculate_size_limit(self, path, size_limit):
494         if size_limit is None:
495             return None
496         else:
497             available_size = self._available_size.get(path)
498             if available_size is None:
499                 size = get_pixbuf_size(path)
500                 if size is None:
501                     # Error reading image size.
502                     return None
503                 available_size = Rectangle(*size)
504                 self._available_size[path] = available_size
505             if available_size.fits_within(size_limit):
506                 return None
507             else:
508                 return size_limit
509
510     def _load_loop(self):
511         while True:
512             request = None
513             self._load_loop_reeval = False
514             found_a_load = False
515             while len(self._load_queue) > 0:
516                 if self._load_queue[0].is_finished():
517                     # The load request has finished, so remove it.
518                     del self._load_queue[0]
519                 else:
520                     request = self._load_queue[0]
521                     found_a_load = True
522                     break
523             if request is None:
524                 # No loads. Find a preload, if any.
525                 for req in self._request_queue.itervalues():
526                     if not req.is_finished():
527                         request = req
528                         break
529             if request is None:
530                 # The request queue needs to be pruned now since there
531                 # may be old load requests that now are finished but
532                 # could not be removed before.
533                 self._prune_queue()
534
535                 # No loads or preloads left. Pause until
536                 # self._load_thread.start() is called.
537                 if self._debug_level > 1:
538                     print
539                     print "LOAD LOOP FINISHED:"
540                     self._print_state(False)
541                 self._load_thread.stop()
542                 yield True
543             else:
544                 if self._debug_level > 1:
545                     print "LOAD LOOP SELECTED TO LOAD", request._path
546                 if found_a_load:
547                     priority = gobject.PRIORITY_HIGH_IDLE
548                 else:
549                     priority = gobject.PRIORITY_DEFAULT_IDLE
550                 self._load_thread.set_priority(priority)
551                 while not (self._load_loop_reeval or request.is_finished()):
552                     loaded_pixels = request.load_some_more()
553                     self._pixels_in_cache += loaded_pixels
554                     self._prune_queue()
555                     yield True
556             # Now it's time to reevaluate which request to work on.
557
558     def _print_state(self, verbose=False):
559         print "------------------------------------------------------------------"
560         print "Pixel limit:", self._pixel_limit
561         print "Cache directory:", self._cache_directory
562         print "Pixels in cache:", self._pixels_in_cache
563         print "Request queue:"
564         for (key, request) in self._request_queue.iteritems():
565             print "    Key:", key
566             print "    Request:"
567             request._print_state(verbose)
568         print "Load queue:"
569         for request in self._load_queue:
570             print "    Request:"
571             request._print_state(verbose)
572         print "Path to requests:"
573         for (path, requests) in self._path_to_requests.iteritems():
574             print "    Path:", path
575             print "    Requests:", requests
576
577     def _prune_queue(self):
578         if self._pixels_in_cache <= self._pixel_limit:
579             return
580         requests_to_prune = []
581         load_requests = set(self._load_queue)
582         pixels_after_pruning = self._pixels_in_cache
583         for (key, request) in self._request_queue.reviteritems():
584             if pixels_after_pruning <= self._pixel_limit:
585                 # Enough pruning.
586                 break
587             if request not in load_requests:
588                 requests_to_prune.append((key, request))
589                 pixels_after_pruning -= request.get_number_of_loaded_pixels()
590         if len(requests_to_prune) > 0:
591             self._load_loop_reeval = True
592             for (key, request) in requests_to_prune:
593                 if self._debug_level > 1:
594                     print "PRUNING", key
595                 self._remove_request(key)
596             gc.collect()
597
598     def _remove_erroneous_request(self, key):
599         self._remove_request(key)
600
601     def _remove_request(self, key):
602         assert key in self._request_queue
603         request = self._request_queue[key]
604         self._pixels_in_cache -= request.get_number_of_loaded_pixels()
605         request.cancel()
606         try:
607             self._load_queue.remove(request)
608         except ValueError:
609             # OK, it wasn't a load.
610             pass
611         del self._request_queue[key]
612         path = key[0]
613         request_set = self._path_to_requests[path]
614         request_set.remove(request)
615         if len(request_set) == 0:
616             del self._path_to_requests[path]
617
618 ######################################################################
619
620 def main(argv):
621     import gtk
622
623     def pixbuf_loaded_cb(pixbuf, full_size):
624         print "Received pixbuf of %dx%d pixels (full size: %dx%d)." % (
625             pixbuf.get_width(), pixbuf.get_height(), full_size[0], full_size[1])
626     def pixbuf_error_cb():
627         print "Error while loading pixbuf."
628
629     loader = CachingPixbufLoader()
630     loader._debug_level = 2
631     loader.set_pixel_limit(10000000000)
632     print "INITIAL:"
633     loader._print_state(False)
634
635     for path in argv[1:]:
636         loader.load(path, (300, 200), pixbuf_loaded_cb, pixbuf_error_cb)
637
638     print
639     print "BEFORE LOAD LOOP START:"
640     loader._print_state(False)
641
642     gtk.main()
643
644 if __name__ == "__main__":
645     import sys
646     main(sys.argv)