Some changes to be more compatible with Python 3 in the future
[joel/kofoto.git] / src / packages / kofoto / imagecache.py
1 """Implementation of the ImageCache class."""
2
3 __all__ = ["ImageCache"]
4
5 import os
6 import Image as PILImage
7 from kofoto.rectangle import Rectangle
8
9 class ImageCache:
10     """A class representing the Kofoto image cache."""
11
12     def __init__(self, cacheLocation, useOrientation=False):
13         """Constructor.
14
15         cachelocation specifies the image cache directory. If
16         useOrientation is true, the image will be rotated according to
17         the orientation attribute.
18         """
19         self.cacheLocation = cacheLocation
20         self.useOrientation = useOrientation
21
22
23     def cleanup(self):
24         """Clean up the cache.
25
26         All cached images whose original images no longer exist in the
27         filesystem will be removed.
28         """
29
30         for dirpath, dirnames, filenames in os.walk(self.cacheLocation,
31                                                     topdown=False):
32             realdir = dirpath[len(self.cacheLocation):]
33             for filename in filenames:
34                 a = os.path.splitext(filename)[0].split("-")
35                 if len(a) >= 4:
36                     realfilename = "-".join(a[0:-3])
37                     mtime = int(a[-1])
38                     try:
39                         currentmtime = os.path.getmtime(
40                             os.path.join(realdir, realfilename))
41                         if currentmtime == mtime:
42                             # Keep.
43                             continue
44                     except OSError:
45                         pass
46                 os.unlink(os.path.join(dirpath, filename))
47             for dirname in dirnames:
48                 # Remove directories if they are empty.
49                 try:
50                     os.rmdir(os.path.join(dirpath, dirname))
51                 except OSError:
52                     pass
53
54
55     def get(self, imageversionOrLocation, widthlimit, heightlimit):
56         """Get a file path to a cached image and the cached image's
57         size.
58
59         imageversionOrLocation could either be an
60         kofoto.shelf.ImageVersion instance or a location string. If
61         it's a location, the path should preferably be normalized
62         (e.g. with os.path.realpath()). If the image does not exist at
63         the given location or cannot be parsed, OSError is raised.
64
65         If the original image doesn't fit within the limits, a smaller
66         version of the image will be created and its path returned. If
67         the original image fits within the limits, a path to a copy of
68         the original image will be returned.
69
70         Returns a tuple of file path, width and height.
71         """
72         if isinstance(imageversionOrLocation, basestring):
73             location = imageversionOrLocation
74             mtime = os.path.getmtime(location)
75             width, height = PILImage.open(location).size
76             orientation = "up"
77         else:
78             imageversion = imageversionOrLocation
79             image = imageversion.getImage()
80             location = imageversion.getLocation()
81             mtime = imageversion.getModificationTime()
82             width, height = imageversion.getSize()
83             if self.useOrientation:
84                 orientation = image.getAttribute(u"orientation")
85                 if not orientation:
86                     orientation = "up"
87             else:
88                 orientation = "up"
89         return self._get(
90             location, mtime, width, height, widthlimit, heightlimit,
91             orientation)
92
93
94     def _get(self, location, mtime, width, height, widthlimit,
95              heightlimit, orientation):
96         """Internal helper method."""
97         # Scale image to fit within limits.
98         w, h = tuple(
99             Rectangle(width, height).downscaled_to(
100                 Rectangle(widthlimit, heightlimit)))
101         if orientation in ["left", "right"]:
102             w, h = h, w
103
104         # Check whether a cached version already exists.
105         path = self._getCachedImagePath(location, mtime, w, h, orientation)
106         if os.path.exists(path):
107             return path, w, h
108
109         # No version of the wanted size existed in the cache. Create
110         # one.
111         directory, _ = os.path.split(path)
112         if not os.path.isdir(directory):
113             os.makedirs(directory)
114         pilimg = PILImage.open(location)
115         if not pilimg.mode in ("L", "RGB", "CMYK"):
116             pilimg = pilimg.convert("RGB")
117         if width > widthlimit or height > heightlimit:
118             if self.useOrientation and orientation in ("left", "right"):
119                 coord = h, w
120             else:
121                 coord = w, h
122             pilimg.thumbnail(coord, PILImage.ANTIALIAS)
123         if self.useOrientation:
124             if orientation == "right":
125                 pilimg = pilimg.rotate(90)
126             elif orientation == "down":
127                 pilimg = pilimg.rotate(180)
128             elif orientation == "left":
129                 pilimg = pilimg.rotate(270)
130         pilimg.save(path, "JPEG")
131         return path, w, h
132
133
134     def _getCachedImagePath(self, location, mtime, width, height,
135                             orientation):
136         """Internal helper method."""
137         drive, drivelessPath = os.path.splitdrive(location)
138         directory, filename = os.path.split(drivelessPath[1:])
139         if drive:
140             directory = os.path.join(drive[0], directory)
141         genname = "%s-%dx%d-%s-%s.jpg" % (
142             filename,
143             width,
144             height,
145             orientation,
146             mtime)
147         return os.path.join(self.cacheLocation, directory, genname)