path: root/trunk/dimbola/utils.py.~1~

# Copyright (C) 2009  Lars Wirzenius <liw@liw.fi>
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


import hashlib
import logging
import math
import os
import StringIO
import subprocess
import tempfile

import gio
import gtk


def abswalk(*args, **kwargs):
    '''Like os.walk, but return absolute pathnames for dirs, filenames.
    
    Arguments are as for os.walk.
    
    For example, abswalk might return the tuple
    ('/etc', ['/etc/default'], ['/etc/passwd', '/etc/group']) where
    os.walk might return ('/etc', ['default'], ['passwd', 'group']).
    
    abswalk os convenient when the caller wants to handle full pathnames
    anyway, as it saves the caller from having to do os.path.join itself.
    
    '''
    def abs(dirname, list):
        return [os.path.join(dirname, x) for x in list]
    for dirname, dirnames, filenames in os.walk(*args, **kwargs):
        yield dirname, abs(dirname, dirnames), abs(dirname, filenames)


def filterabswalk(is_ok, *args, **kwargs):
    '''Like abswalk, but filenames (not dirnames) can be filtered.
    
    The is_ok argument is a function that gets the fully qualified name of
    a file (not directory) and returns True/False to indicate whether it
    should be included in the results.
    
    All other arguments are as for os.walk.
    
    '''
    for dirname, dirnames, pathnames in abswalk(*args, **kwargs):
        filenames = [x for x in pathnames if is_ok(x)]
        yield dirname, dirnames, filenames
        

def safe_copy(input_name, output_name, callback):
    """Copy contents of input_name to new file called output_name.
    
    If the output_name already exists, fail. If anything else goes
    wrong, fail. Ensure the data is on disk using fsync on the output
    name and on the directory containing the output.
    
    The permissions and other stat information for the input are NOT
    copied to the output.
    
    """
    
    infd = os.open(input_name, os.O_RDONLY)
    outfd = os.open(output_name, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0644)
    
    total_copied = 0
    while True:
        data = os.read(infd, 1024**2)
        if not data:
            break
        os.write(outfd, data)
        total_copied += len(data)
        if callback:
            callback(input_name, output_name, total_copied)

    os.close(infd)
    os.fsync(outfd)
    os.close(outfd)
    
    output_dir = os.path.dirname(output_name) or "."
    dirfd = os.open(output_dir, os.O_RDONLY)
    os.fsync(dirfd)
    os.close(dirfd)


def filter_cmd(argv, input_data):
    '''Filter input data through an external command.'''

    fd, name = tempfile.mkstemp()
    os.write(fd, input_data)
    os.lseek(fd, 0, 0)
    os.remove(name)

    p = subprocess.Popen(argv, stdin=fd, stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE)
    stdout, stderr = p.communicate()
    if p.returncode:
        raise Exception('command %s failed: exit code %s\n%s' %
                        (argv, p.returncode, stderr or ''))
    return stdout


def image_data_to_pixbuf(image_data):
    '''Create a gdk.Pixbuf out of some image data.'''
    
    loader = gtk.gdk.PixbufLoader()
    loader.write(image_data)
    loader.close()
    return loader.get_pixbuf()


def pixbuf_to_image_data(pixbuf, format, options=None): # pragma: no cover
    f = StringIO.StringIO()
    def save_func(buf):
        f.write(buf)
        return True
    pixbuf.save_to_callback(save_func, format, options=options)
    return f.getvalue()
    
    
def image_data_to_image_data(data, format, options=None): # pragma: no cover
    pixbuf = image_data_to_pixbuf(data)
    return pixbuf_to_image_data(pixbuf, format, options=options)
    

def scale_pixbuf(pixbuf, maxw, maxh):
    '''Scale a pixbuf so it fits within maxw and maxh.
    
    Keep aspect ratio.
    
    '''
    
    w = pixbuf.get_width()
    h = pixbuf.get_height()
    
    fw = float(maxw) / float(w)
    fh = float(maxh) / float(h)
    f = min(fw, fh)
    w2 = int(f * w)
    h2 = int(f * h)
    assert w2 <= maxw
    assert h2 <= maxh
    
    return pixbuf.scale_simple(w2, h2, gtk.gdk.INTERP_BILINEAR)


def rotate_pixbuf(pixbuf, angle):
    '''Rotate pixbuf in 90 degree angles.'''
    values = {
        90: gtk.gdk.PIXBUF_ROTATE_COUNTERCLOCKWISE,
        180: gtk.gdk.PIXBUF_ROTATE_UPSIDEDOWN,
        270: gtk.gdk.PIXBUF_ROTATE_CLOCKWISE,
    }
    return pixbuf.rotate_simple(values.get(angle, gtk.gdk.PIXBUF_ROTATE_NONE))


def encode_dnd_tagids(tagids):
    '''Encode a list of tagids for drag-and-drop.'''
    return ' '. join(str(tagid) for tagid in tagids)


def decode_dnd_tagids(encoded):
    '''Reverse operation of encode_dnd_tagids.'''
    return [int(s) for s in encoded.split()]


class TreeBuilder(object):

    '''Build a tree out of a sequence of nodes.
    
    The caller provides zero or more node descriptions (see add method).
    After the caller is done (see done method), they can access the
    tree built out of the nodes as the tree property. The tree is
    represented as a list of node tuples (nodeid, data, child_nodes).
    
    If a node should have a parent, but the parent is not added, the
    node becomes a root node.
    
    '''
    
    def __init__(self):
        # We store nodes in a dictionary indexed by the nodeid.
        # The value is a tuple of (data, parentid, childids).
        # Childids is initially an empty set, it'll be used by
        # the done method.
        self.nodes = dict()

    def add(self, nodeid, data, sortkey, parentid):
        '''Add a node to tree.
        
        nodeid is the identifier of the node itself.
        sortkey is used when sorting children with the same parent.
        parentid is the identifier of its parent, or None.
        
        data is the data associated with the node.
        
        '''
        
        self.nodes[nodeid] = (data, sortkey, parentid, set())

    def done(self):
        '''Caller is done adding nodes, compute the tree.
        
        Caller MUST call this; until this is called, self.tree does not
        exist.
        
        '''

        # First we put each node into its parents' childids.
        # If parent is missing, we pretend it was always None.
        for nodeid in self.nodes:
            data, sortkey, parentid, children = self.nodes[nodeid]
            if parentid in self.nodes:
                self.nodes[parentid][3].add(nodeid)
            else:
                self.nodes[nodeid] = (data, sortkey, None, children)

        # Next we find all root nodes: all nodes whose parentid is None.
        roots = [nodeid 
                 for nodeid in self.nodes 
                 if self.nodes[nodeid][2] is None]

        # Next we build the tree for each root node, and add those
        # to the tree.
        rootlist = [(self.nodes[rootid][1], rootid) for rootid in roots]
        rootlist.sort()
        roots = [rootid for sortkey, rootid in rootlist]
        self.tree = [self.build_one_tree(rootid) for rootid in roots]

    def build_one_tree(self, rootid):
        data, sortkey, parentid, childids = self.nodes[rootid]
        childlist = [(self.nodes[kid][1], kid) for kid in childids]
        childlist.sort()
        childids = [kid for sortkey, kid in childlist]
        return (rootid, data, 
                [self.build_one_tree(childid) for childid in childids])


class DcrawTypeCache(object):

    '''Cache 'dcraw -i' results.
    
    dcraw does not export a list of MIME types it recognizes, but it does
    have an option to test whether it supports the format of a particular
    file. That's slow, so we cache the results using this class.
    
    The results are stored in a format compatible with what
    gtk.gdk.pixbuf_get_formats returns: a list of dictionaries with
    keys 'name', 'mime_types', and 'extension'. (This is a subset of
    the keys for pixbufs.)

    '''
    
    def __init__(self):
        self.formats = []
        self.fail_extensions = set()
        self.fail_mime_types = set()

    def update(self, format, mime_type, extension):
        if mime_type not in format['mime_types']:
            format['mime_types'].append(mime_type)
        for ext in [extension, extension.lower(), extension.upper()]:
            if ext not in format['extensions']:
                format['extensions'].append(ext)

    def add_format(self, name, mime_type, extension):
        for format in self.formats:
            if format['name'] == name:
                self.update(format, mime_type, extension)
                return
            elif mime_type in format['mime_types']:
                self.update(format, mime_type, extension)
                return
            elif extension in format['extensions']:
                self.update(format, mime_type, extension)
                return

        format = {
            'name': name,
            'mime_types': [],
            'extensions': [],
        }
        self.update(format, mime_type, extension)
        self.formats.append(format)

    def extension_is_known(self, ext):
        return [x for x in self.formats if ext in x['extensions']]

    def mime_type_is_known(self, mimetype):
        return [x for x in self.formats if mimetype in x['mime_types']]
        
    def get_mime_type(self, filename):
        f = gio.File(path=filename)
        fi = f.query_info(gio.FILE_ATTRIBUTE_STANDARD_CONTENT_TYPE)
        return gio.content_type_get_mime_type(fi.get_content_type())

    def get_dcraw(self, filename):
        try:
            p = subprocess.Popen(['dcraw', '-i', filename], 
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE)
            stdout, stderr = p.communicate('')
        except OSError, e: # pragma: no cover
            logging.debug('Cannot run dcraw: %s' % str(e))
            return None
        if p.returncode == 0:
            prefix = '%s is a ' % filename
            if stdout.startswith(prefix):
                stdout = stdout[len(prefix):]
            suffix = ' image.\n'
            if stdout.endswith(suffix):
                stdout = stdout[:-len(suffix)]
            return stdout
        else:
            return None

    def add_from_file(self, filename):
        prefix, ext = os.path.splitext(filename)
        if ext.startswith('.'):
            ext = ext[1:]
        if self.extension_is_known(ext):
            return
        if ext in self.fail_extensions:
            return

        mime_type = self.get_mime_type(filename)
        if self.mime_type_is_known(mime_type):
            return
        if mime_type in self.fail_mime_types:
            return
            
        desc = self.get_dcraw(filename)
        if desc is None:
            self.fail_extensions.add(ext)
            self.fail_mime_types.add(mime_type)
        else:
            self.add_format(ext, mime_type, ext)

    def supported(self, filename):
        self.add_from_file(filename)
        name, ext = os.path.splitext(filename)
        if ext.startswith('.'):
            ext = ext[1:]
        if self.extension_is_known(ext):
            return True
        mime_type = self.get_mime_type(filename)
        return self.mime_type_is_known(mime_type)
        
    
def draw_star(drawable, gc, x, y, dim): # pragma: no cover
    '''Draw a five-pointed star.
    
    The star will be drawn inside a square of dim pixels, whose top left
    corner is at (x,y). The star will be filled. The graphics context gc
    is used for drawing and filling.
    
    '''
    
    # To follow this code, imagine a circle inscribed in the square.
    # The star is a pentagram drawn inside the circle, situated so that
    # one of its points is pointing upwards. The five points are called
    # A through E. Inside the pentagram is an upside down pentagon. It's
    # lowest point is directly below A (same x co-ordinate), and is called
    # F. We draw the pentagram by drawing three filled triangles: ACF,
    # ADF, and BEF.
    #
    # The co-ordinates of the six points are a bit tricky, or I am stupid.
    # First we find the co-ordinates with the assumption that the center of
    # the square (and circle and pentagram) is at origin, then we displace
    # them to the right place. Note also that screen and geometrical
    # y-axis are in opposite direction.
    #
    # The radius of the circle is R = dim/2.
    #
    # The angle AOB is 2*pi/5.
    #
    # A is simple: (0, R).
    # 
    # B: angle between FB and x-axis is (pi/2 - AOB) = (pi/2 - 2*pi/5) =
    # (5*pi/10 - 4*pi/10) = pi/10 = alpha. 
    # Thus B = (R*cos alpha, R*sin alpha).
    #
    # C: angle between FC and x-axis is (BOC - alpha) = (2*pi/5 - alpha) =
    # (2*pi/5 - pi/10) = (4*pi/10 - pi/10) = 3*pi/10 = beta.
    # Thus C = (R*cos beta, R*sin beta).
    #
    # D = (-Cx, Cy).
    #
    # E = (-Bx, By).
    #
    # F: Let P = (Cx, By), Z = (0, By). The triangle EZF is shaped like
    # EPC, but smaller. EZ/ZF = EP/PC <=> ZF = EZ*PC/EP. Also,
    # ZF = ZO + OF so OF = EZ*PC/EP - ZO. We have the co-ordinates for
    # everything except F, and F = (0, OF). Thus: 
    # Fy = -(Bx*(By-Cy)/(Bx+Cx) - By) = Bx*(Cy-By)/(Bx+Cx)+By.
    
    R = float(dim) / 2.0
    alpha = math.pi / 10.0
    beta = -3.0 * math.pi / 10.0
    
    # These calculations are done in normal math co-ordinate system.
    # (Y grows upwards.)
    A = (0, R)
    B = (R * math.cos(alpha), R * math.sin(alpha))
    C = (R * math.cos(beta), R * math.sin(beta))
    D = (-C[0], C[1])
    E = (-B[0], B[1])
    F = (0, -(B[0] * (B[1] - C[1]) / (B[0] + C[0]) - B[1]))
    F = (0, B[0] * (C[1] - B[1]) / (B[0] + C[0]) + B[1])

    # Transform co-ordinates to screen: move origin to center of square,
    # and change direction of Y axis.
    def xform(coords):
        return int(x + R + coords[0]), int(y + R - coords[1])

    A = xform(A)
    B = xform(B)
    C = xform(C)
    D = xform(D)
    E = xform(E)
    F = xform(F)
    
    # Draw the three triangles.
    drawable.draw_polygon(gc, True, (A, F, C, A))
    drawable.draw_polygon(gc, True, (A, F, D, A))
    drawable.draw_polygon(gc, True, (B, E, F, B))


def draw_stars(n_stars, drawable, gc, x, y, dim): # pragma: no cover
    '''Like draw_star, but draws n_stars stars.
    
    The drawable MUST be wide enough to have space for five (5) stars.
    That area will be cleared.
    
    '''

    drawable.clear_area(x, y, 5 * dim, dim)
    for i in range(n_stars):
        draw_star(drawable, gc, x + i*dim, y, dim)


def sha1(filename): # pragma: no cover
    '''Compute SHA1 checksum of a file.
    
    Return None if there were errors.
    
    '''
    try:
        f = file(filename)
    except IOError:
        return None
    
    c = hashlib.new('sha1')
    while True:
        data = f.read(64*1024)
        if not data:
            break
        c.update(data)
    f.close()
    return c.hexdigest()