diff --git a/gpkit/keydict.py b/gpkit/keydict.py index 27f252902..7285f03b4 100644 --- a/gpkit/keydict.py +++ b/gpkit/keydict.py @@ -44,14 +44,13 @@ class KeyMap: attribute, that key can be set and accessed normally. """ collapse_arrays = False - keymap = [] + keymap = defaultdict(set) log_gets = False varkeys = None vks = None def __init__(self, *args, **kwargs): "Passes through to super().__init__ via the `update()` method" - self.keymap = defaultdict(set) self._unmapped_keys = set() self.owned = set() self.update(*args, **kwargs) # pylint: disable=no-member @@ -106,7 +105,7 @@ def __contains__(self, key): # pylint:disable=too-many-return-statements " for value %s" % (key, idx, super().__getitem__(key))) # pylint: disable=no-member return True - return key in self.keymap + return any(super().__contains__(k) for k in self.keymap[key]) def update_keymap(self): "Updates the keymap with the keys in _unmapped_keys" @@ -172,7 +171,10 @@ def __call__(self, key): # if uniting is ever a speed hit, cache it def __getitem__(self, key): "Overloads __getitem__ and [] access to work with all keys" key, idx = self.parse_and_index(key) - keys = self.keymap[key] + if super().__contains__(key): + keys = [key] + else: + keys = self.keymap[key] if not keys: del self.keymap[key] # remove blank entry added by defaultdict raise KeyError(key)