diff --git a/tableaudocumentapi/multilookup_dict.py b/tableaudocumentapi/multilookup_dict.py index 64b742a..6677b07 100644 --- a/tableaudocumentapi/multilookup_dict.py +++ b/tableaudocumentapi/multilookup_dict.py @@ -39,15 +39,25 @@ def _populate_indexes(self): self._indexes['alias'] = _build_index('alias', self) self._indexes['caption'] = _build_index('caption', self) + def _get_real_key(self, key): + if key in self._indexes['alias']: + return self._indexes['alias'][key] + if key in self._indexes['caption']: + return self._indexes['caption'][key] + + return key + def __setitem__(self, key, value): + real_key = self._get_real_key(key) + alias = _resolve_value('alias', value) caption = _resolve_value('caption', value) if alias is not None: - self._indexes['alias'][alias] = key + self._indexes['alias'][alias] = real_key if caption is not None: - self._indexes['caption'][caption] = key + self._indexes['caption'][caption] = real_key - dict.__setitem__(self, key, value) + dict.__setitem__(self, real_key, value) def get(self, key, default_value=_no_default_value): try: @@ -58,9 +68,5 @@ def get(self, key, default_value=_no_default_value): raise def __getitem__(self, key): - if key in self._indexes['alias']: - key = self._indexes['alias'][key] - elif key in self._indexes['caption']: - key = self._indexes['caption'][key] - - return dict.__getitem__(self, key) + real_key = self._get_real_key(key) + return dict.__getitem__(self, real_key) diff --git a/test/test_multidict.py b/test/test_multidict.py index 0a78e9d..4544c49 100644 --- a/test/test_multidict.py +++ b/test/test_multidict.py @@ -61,3 +61,12 @@ def test_multilookupdict_get_returns_default_value(self): def test_multilookupdict_get_returns_value(self): actual = self.mld.get('baz') self.assertEqual(1, actual['value']) + + def test_multilookupdict_can_set_item(self): + before = self.mld['baz'] + self.mld['baz'] = 4 + self.assertEqual(4, self.mld['baz']) + + def test_multilookupdict_can_set_new_item(self): + self.mld['wakka'] = 1 + self.assertEqual(1, self.mld['wakka'])