laydi/fluents/annotations.py

from fluents import dataset

_dim_annotation_handlers = {}

def get_dim_annotations(dimname, annotation, ids):
    """Returns a list of annotations corresponding to the given ids in 
    dimension dimname"""
    global _dim_annotation_handlers

    if _dim_annotation_handlers.has_key(dimname):
        return _dim_annotation_handlers[dimname].get_annotations(annotation, ids)
    return None

def set_dim_handler(dimname, handler):
    """Set the handler for the given dimension."""
    global _dim_annotation_handlers
    _dim_annotation_handlers[dimname] = handler
    
def get_dim_handler(dimname):
    """Get the handler for the given dimension."""
    global _dim_annotation_handlers
    if _dim_annotation_handlers.has_key(dimname):
        return _dim_annotation_handlers[dimname]
    else: 
        return None


class AnnotationHandler:
    def __init__(self):
        pass

    def get_annotations(self, annotationname, ids, default=None):
        return None

    def get_annotation_names(self):
        return []


class DictAnnotationHandler(AnnotationHandler):

    def __init__(self, d={}):
        self._dict = d
    
    def get_annotations(self, annotationname, ids, default=None):
        d = self._dict
        retval = []
        for id in ids:
            if d[annotationname].has_key(id):
                retval.append(d[annotationname][id])
            else:
                retval.append(default)
        return retval
    
    def add_annotations(self, annotationname, d):
        self._dict[annotationname] = d

    def get_annotation_names(self):
        return self._dict.keys()
        

def read_annotations_file(filename):
    """Read annotations from file.
    
    Reads annotations from a tab delimited file of the format::
     dimname     annotation_name1    annotation_name2 ...
     id1         Foo                 0.43
     id2         Bar                 0.59
    """
    
    ann = DictAnnotationHandler()
    dimname = None
    annotation_dicts = []
    annotation_names = []
    
    fd = open(filename)

    ## Read the first line, which contains the dimension name and
    ## annotation names.
    line = fd.readline()
    values = [x.strip() for x in line.split('\t')]
    dimname = values[0]  
    annotation_names = values[1:]
    annotation_dicts = [{} for x in annotation_names]

    ## Read the lines containing the annotations. The first value on
    ## each line is an id along the dimension.
    while line:
        values = [x.strip() for x in line.split('\t')]
        for i, x in enumerate(values[1:]):
            annotation_dicts[i][values[0]] = x
        line = fd.readline()

    fd.close()

    ## Add everything to the annotation object and add the object to
    ## the specified dimension.
    for i, a in enumerate(annotation_names):
        ann.add_annotations(a, annotation_dicts[i])

    _dim_annotation_handlers[dimname] = ann
    return ann