refactor for speed

2 files changed, 15 insertions, 14 deletions

diff --git a/meliaereader/reader.py b/meliaereader/reader.py
index a0a2f4eb..08b84887 100644
--- a/meliaereader/reader.py
+++ b/meliaereader/reader.py
@@ -22,23 +22,25 @@ import json
 class MeliaeReader(object):
 
     def __init__(self):
-        self._objs = []
+        self._objs = {}
 
     def __iter__(self):
-        return iter(self._objs)
+        return iter(self._objs.values())
 
     def __contains__(self, ref):
-        return any(o['address'] == ref for o in self)
+        return ref in self._objs
 
     def __len__(self):
         return len(self._objs)
 
     def read(self, filename):
         with open(filename) as f:
-            self._objs.extend(json.loads(line) for line in f)
+            for line in f:
+                obj = json.loads(line)
+                self._objs[obj['address']] = obj
 
     def get_total_size(self):
-        return self.get_size(self._objs)
+        return self.get_size(self._objs.values())
 
     def get_size(self, objs):
         return sum(o['size'] for o in objs)
@@ -63,9 +65,8 @@ class MeliaeReader(object):
         return closure
 
     def get_object(self, ref):
-        for obj in self:
-            if obj['address'] == ref:
-                return obj
+        if ref in self:
+            return self._objs[ref]
         raise Exception('No object with address {}'.format(ref))
 
     def get_closure_of_type(self, typename):
diff --git a/meliaereader/reader_tests.py b/meliaereader/reader_tests.py
index c50e97c1..4ef80483 100644
--- a/meliaereader/reader_tests.py
+++ b/meliaereader/reader_tests.py
@@ -54,7 +54,7 @@ class MeliaeReaderTests(unittest.TestCase):
         self.assertEqual(len(mr), 0)
 
     def test_reads_file_with_one_object(self):
-        filename = self.make_file(self.make_object())
+        filename = self.make_file(self.make_object(address=1))
         mr = meliaereader.MeliaeReader()
         mr.read(filename)
         self.assertEqual(len(mr), 1)
@@ -85,22 +85,22 @@ class MeliaeReaderTests(unittest.TestCase):
         self.assertEqual(mr.get_size([obj]), 42)
 
     def test_readings_adds_to_existing_objs(self):
-        filename_1 = self.make_file(self.make_object())
-        filename_2 = self.make_file(self.make_object())
+        filename_1 = self.make_file(self.make_object(address=1))
+        filename_2 = self.make_file(self.make_object(address=2))
         mr = meliaereader.MeliaeReader()
         mr.read(filename_1)
         mr.read(filename_2)
         self.assertEqual(len(mr), 2)
 
     def test_reports_object_types(self):
-        filename = self.make_file(self.make_object(type='foo'))
+        filename = self.make_file(self.make_object(type='foo', address=1))
         mr = meliaereader.MeliaeReader()
         mr.read(filename)
         self.assertEqual(mr.get_types(), set(['foo']))
 
     def test_reports_objs_of_type(self):
-        obj_1 = self.make_object(type='foo')
-        obj_2 = self.make_object(type='bar')
+        obj_1 = self.make_object(type='foo', address=1)
+        obj_2 = self.make_object(type='bar', address=2)
         filename = self.make_file(obj_1, obj_2)
         mr = meliaereader.MeliaeReader()
         mr.read(filename)