terom@20: import shelve
terom@20: 
terom@20: class TagDB (object) :
terom@20:     _imgs_cache = {}
terom@20: 
terom@20:     def __init__ (self, read_only=True) :
terom@20:         self.img_tags = shelve.open("img_tags", read_only and 'r' or 'c')
terom@20:         self.tag_imgs = shelve.open("tag_imgs", read_only and 'r' or 'c')
terom@20: 
terom@20:     def tag (self, img, tag) :
terom@20:         """
terom@20:             Associate the given image with the given tag
terom@20:         """
terom@20: 
terom@20:         if img not in self.img_tags :
terom@20:             self.img_tags[img] = set([tag])
terom@20:         else :
terom@20:             s = self.img_tags[img]
terom@20:             s.add(tag)
terom@20:             self.img_tags[img] = s
terom@20: 
terom@20: 
terom@20:         if tag not in self.tag_imgs :
terom@20:             self.tag_imgs[tag] = set([img])
terom@20:         else :
terom@20:             s = self.tag_imgs[tag]
terom@20:             s.add(img)
terom@20:             self.tag_imgs[tag] = s
terom@20: 
terom@20:         print "%s <-> %s" % (img, tag)
terom@20:     
terom@20:     def imgs (self, tags) :
terom@20:         """
terom@20:             Get the set of images that have the given set of tags
terom@20:         """
terom@20: 
terom@20:         cache_key = "/".join(tags)
terom@20:         
terom@20:         if cache_key in self._imgs_cache :
terom@20:             return self._imgs_cache[cache_key]
terom@20: 
terom@20:         if not tags :
terom@20:             return set(self.img_tags.keys())
terom@20: 
terom@20:         img_sets = [self.tag_imgs[tag] for tag in tags]
terom@20:         
terom@20:         res = None
terom@20: 
terom@20:         for img_set in img_sets :
terom@20:             if res :
terom@20:                 res = res & img_set
terom@20:             else :
terom@20:                 res = img_set
terom@20: 
terom@20:         self._imgs_cache[cache_key] = res
terom@20:         
terom@20:         return res
terom@20:     
terom@20:     def tags (self, tags) :
terom@20:         """
terom@20:             Get the set of tags that are present in the set of images specified by these tags, sorted by count
terom@20: 
terom@20:             This is currently implemented quite inefficiently... giev SQL db?
terom@20:         """
terom@20:         
terom@20:         imgs = self.imgs(tags)
terom@20:         
terom@20:         ret = []
terom@20:         
terom@20:         for tag in self.tag_imgs.keys() :
terom@20:             if tag in tags :
terom@20:                 continue
terom@20: 
terom@20:             count = len(self.tag_imgs[tag] & imgs)
terom@20: 
terom@20:             if count :
terom@20:                 ret.append((tag, count))
terom@20: 
terom@20:         def my_cmp ((at, ac), (bt, bc)) :
terom@20:             return cmp((ac, at), (bc, bt))
terom@20: 
terom@20:         ret.sort(reverse=True)
terom@20: 
terom@20:         return ret
terom@20:     
terom@20:     def close (self) :
terom@20:         self.img_tags.close()
terom@20:         self.tag_imgs.close()
terom@20: