diff --git a/system/annotations.py b/system/annotations.py index eaf5139..17cc4bc 100644 --- a/system/annotations.py +++ b/system/annotations.py @@ -1,33 +1,65 @@ -import logger + +from sets import Set as set +set.update = set.union_update + +class AnnotationsException(Exception): + pass class Annotations: def __init__(self, *dimensions): """ Initializes a new Annotation with the given dimension labels. dimensions is a list of dimension labels. """ - if len(dimensions) == 2: - logger.log('error', 'Annotations only supports two dimensions.') + if len(dimensions) != 2: + msg = 'Annotations only supports two dimensions.' + raise AnnotationsException(msg) self.dimensions = {} for d in dimensions: self.dimensions[d] = {} def add_annotations(self, dim, id, ann_dim, annotations): - if not self.dimensions[dim]: - logger.log('warning', 'Annotations object does not contain dimension %s' % dim) - return None + """ Adds new annotations. + dim: the dimension in which the new data should be added. + id: the identifier that should be annotated. + ann_dim: the dimension of the annotations to id. + annotations: the new annotations to id. + Returns the total set of annotations to id. + """ + if not self.has_dimension(dim): + msg = 'Annotations object does not contain dimension %s' % dim + raise AnnotationsException(msg) + + if not self.has_dimension(ann_dim): + msg = 'Annotations object does not contain dimension %s' % ann_dim + raise AnnotationsException(msg) + for a in annotations: - self.annotations[ann_dim][a].add(id) - self.annotations[dim][id] = annotations + if not self.dimensions[ann_dim].has_key(a): + self.dimensions[ann_dim][a] = set() + self.dimensions[ann_dim][a].add(id) + + if not self.dimensions[dim].has_key(id): + self.dimensions[dim][id] = set() + self.dimensions[dim][id].update(annotations) + + return self.dimensions[dim][id] - def get_annotations(self, dim, id): - if not self.dimensions[dim]: - logger.log('warning', 'Annotations object does not contain dimension %s' % dim) - return None + def get_annotations(self, dim, id, ann_dim): + """Returns all annotations to id. + dim: the dimension where id can be found. + id: the id to retrieve annotations for. + """ + if not self.has_dimension(dim): + msg = 'Annotations object does not contain dimension %s' % dim + raise AnnotationsException(msg) if self.dimensions[dim].has_key(id): - return self.dimensions[dim][id].values() - return None + return self.dimensions[dim][id] + return set() def has_dimension(self, dim): + """ Retuns true if the Annotations object indexes dim. + """ return self.dimensions.has_key(dim) + diff --git a/test/system/annotationtest.py b/test/system/annotationtest.py index 9b85575..b9982b1 100644 --- a/test/system/annotationtest.py +++ b/test/system/annotationtest.py @@ -3,6 +3,7 @@ import unittest import sys sys.path.append('../../system') from annotations import * +from sets import Set as set class AnnotationTest(unittest.TestCase): @@ -11,9 +12,21 @@ class AnnotationTest(unittest.TestCase): def testCreation(self): a = Annotations('x', 'y') + assert a.has_dimension('x') assert a.has_dimension('y') assert not a.has_dimension('z') + def testAddAnnotations(self): + ann = Annotations('genes', 'go') + go = set(['GO:0', 'GO:1']) + ann.add_annotations('genes', 'BadGene', 'go', go) + genes = set(['BadGene']) + + assert ann.get_annotations('genes', 'BadGene', 'go') == go + assert ann.get_annotations('go', 'GO:0', 'genes') == genes + assert ann.get_annotations('go', 'GO:1', 'genes') == genes + assert ann.get_annotations('go', 'GO:2', 'genes') == set() + if __name__ == '__main__': unittest.main()