From 439eb85652613cd5a966920ac383fba920f07e2a Mon Sep 17 00:00:00 2001 From: nocturn9x Date: Mon, 1 Jun 2020 14:39:32 +0000 Subject: [PATCH] Fixed the _contains__ dunder method --- ttlcollections/objects.py | 7 ++++++- ttlcollections/structures.py | 6 +++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/ttlcollections/objects.py b/ttlcollections/objects.py index 5eb5e95..13c875e 100644 --- a/ttlcollections/objects.py +++ b/ttlcollections/objects.py @@ -2,7 +2,7 @@ class TTLItem(object): - """An abstraction layer over a queue item. This class is meant + """An abstraction layer over a TTLed item. This class is meant to be used internally and should not be instantiated directly :param obj: The object to be wrapped @@ -21,3 +21,8 @@ class TTLItem(object): if isinstance(other, TTLItem): other = other.obj return self.obj.__lt__(other) + + def __eq__(self, other): + if isinstance(other, TTLItem): + other = other.obj + return self.obj.__eq__(other) diff --git a/ttlcollections/structures.py b/ttlcollections/structures.py index 6bffba0..46a4254 100644 --- a/ttlcollections/structures.py +++ b/ttlcollections/structures.py @@ -114,7 +114,7 @@ class TTLQueue: def __contains__(self, item): """Implements item in self""" - return self._queue.__contains__(item) + return self._queue.__contains__(TTLitem(item, None)) class TTLStack: """A stack (LIFO) with per-item time to live (TTL) @@ -220,7 +220,7 @@ class TTLStack: def __contains__(self, item): """Implements item in self""" - return self._stack.__contains__(item) + return self._stack.__contains__(TTLItem(item, None)) class TTLHeap(TTLQueue): @@ -272,7 +272,7 @@ class TTLHeap(TTLQueue): def __contains__(self, item): """Implements item in self""" - return self._queue.__contains__(item) + return self._queue.__contains__(TTLitem(item, None)) def put(self, element, ttl: int = 0): """Puts an item onto the queue