From ec7f0797a9c11deb02a1c1e11a66436bcb5acce0 Mon Sep 17 00:00:00 2001 From: Andrei Date: Sun, 27 May 2018 17:07:32 -0700 Subject: [PATCH] Fix the way HashMap.select() works --- disco/util/hashmap.py | 14 +++++--------- tests/util/hashmap.py | 11 +++++++++++ 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/disco/util/hashmap.py b/disco/util/hashmap.py index b689032..4c999b1 100644 --- a/disco/util/hashmap.py +++ b/disco/util/hashmap.py @@ -27,16 +27,12 @@ class HashMap(dict): def find_one(self, predicate): return next(self.find(predicate), None) - def select(self, *args, **kwargs): - if kwargs: - args += tuple(kwargs) - + def select(self, **kwargs): for obj in self.values(): - for check in args: - for k, v in six.iteritems(check): - if getattr(obj, k) != v: - break - yield obj + for k, v in six.iteritems(kwargs): + if getattr(obj, k) != v: + break + yield obj def select_one(self, **kwargs): return next(self.select(**kwargs), None) diff --git a/tests/util/hashmap.py b/tests/util/hashmap.py index 3ffa931..41eacf9 100644 --- a/tests/util/hashmap.py +++ b/tests/util/hashmap.py @@ -44,3 +44,14 @@ def test_hashmap_items(hashmap): assert len(list(hashmap.items())) == 100000 assert list(hashmap.items())[0][0] == 0 assert list(hashmap.items())[0][1] == hashmap[0] + + +def test_hashmap_select_one(): + class Test(object): + x = 1 + y = 2 + hashmap = HashMap() + hashmap['x'] = Test() + + assert hashmap.select_one(x=1) == hashmap['x'] + assert hashmap.select_one(x=2) == None