Overloading __getitem__() for two key types in Python

Question

I've got a collection of users with two types of IDs (a_id, b_id) where one is a positive integer and the other is a six-letter string. For my collection class, I'd like to be able to look up users using either type of ID.

Is there something wrong with using __contains__ and __getitem__ for either type of ID?

class UserList: def __init__(self, users_json): self.a_ids = {} self.b_ids = {} for user in users_json: a_id = user['a_id'] b_id = user['b_id'] self.a_ids[a_id] = user self.b_ids[b_id] = user def __contains__(self, some_id): return some_id in self.a_ids or some_id in self.b_ids def __getitem__(self, some_id): try: return self.a_ids[some_id] except KeyError: return self.b_ids[some_id] 

Update: This is for Python 3.x, and there is no implementation of __setitem__; updating users is handled in separate API functions.

Answer 1

• From Python 3.3, the collections module added collections.ChainMap. Utilizing this class can simplify your code.
• You should probably make self.a_ids and self.b_ids a list. And index to get which one you want, such as self.ids[0]. This list is created when using ChainMap anyway too.

And so I'd change your code to something more like:

from collections import ChainMap class UserList: def __init__(self, users): self._lookup = ChainMap({}, {}) self.ids = self._lookup.maps for user in users: self.ids[0][user['a_id']] = user self.ids[1][user['b_id']] = user def __contains__(self, some_id): return some_id in self._lookup def __getitem__(self, some_id): return self._lookup[some_id] 

Alternately, you could subclass ChainMap, and not need to write any of the methods.

from collections import ChainMap class UserList(ChainMap): def __init__(self, users): super().__init__({}, {}) for user in users: self.maps[0][user['a_id']] = user self.maps[1][user['b_id']] = user 

However, at that point, there's not much benefit to using a class. And so I'd just use a raw ChainMap.