From 4bdeef05f6ad59ab18fb5322ee4a30e3dd29e4c0 Mon Sep 17 00:00:00 2001
From: einarr <einarr@pvv.ntnu.no>
Date: Sun, 16 Apr 2006 23:30:01 +0000
Subject: [PATCH]  * annotations.py: The Annotations class now work. Querying
 along any    axis is supported, but only two dimensional Annotations can be
 created.

---
 system/annotations.py         | 60 +++++++++++++++++++++++++++--------
 test/system/annotationtest.py | 13 ++++++++
 2 files changed, 59 insertions(+), 14 deletions(-)

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()