add dict-style lookup on rId to RelationshpCollctn

Steve Canny · Steve Canny · commit 425ce84a8e72 · 2013-08-05T15:08:40.000-07:00
diff --git a/opc/package.py b/opc/package.py
@@ -140,9 +140,18 @@ def __init__(self, baseURI):
         self._baseURI = baseURI
         self._rels = []
 
-    def __getitem__(self, idx):
-        """Implements access by subscript, e.g. rels[9]"""
-        return self._rels.__getitem__(idx)
+    def __getitem__(self, key):
+        """
+        Implements access by subscript, e.g. ``rels[9]``. It also implements
+        dict-style lookup of a relationship by rId, e.g. ``rels['rId1']``.
+        """
+        if isinstance(key, basestring):
+            for rel in self._rels:
+                if rel.rId == key:
+                    return rel
+            raise KeyError("no rId '%s' in RelationshipCollection" % key)
+        else:
+            return self._rels.__getitem__(key)
 
     def __len__(self):
         """Implements len() built-in on this object"""
diff --git a/tests/test_package.py b/tests/test_package.py
@@ -152,6 +152,19 @@ def it_supports_indexed_access(self):
         except IndexError:
             pass
 
+    def it_has_dict_style_lookup_of_rel_by_rId(self):
+        rel = Mock(name='rel', rId='foobar')
+        rels = RelationshipCollection(None)
+        rels._rels.append(rel)
+        assert rels['foobar'] == rel
+
+    def it_should_raise_on_failed_lookup_by_rId(self):
+        rel = Mock(name='rel', rId='foobar')
+        rels = RelationshipCollection(None)
+        rels._rels.append(rel)
+        with pytest.raises(KeyError):
+            rels['barfoo']
+
     def it_can_add_a_relationship(self, _Relationship_):
         baseURI, rId, reltype, target, external = (
             'baseURI', 'rId9', 'reltype', 'target', False
