RustPython
diff --git a/‎Lib/queue.py
Lines changed: 59 additions & 1 deletion b/‎Lib/queue.py
Lines changed: 59 additions & 1 deletion
@@ -10,7 +10,15 @@
 except ImportError:
     SimpleQueue = None
 
-__all__ = ['Empty', 'Full', 'Queue', 'PriorityQueue', 'LifoQueue', 'SimpleQueue']
+__all__ = [
+    'Empty',
+    'Full',
+    'ShutDown',
+    'Queue',
+    'PriorityQueue',
+    'LifoQueue',
+    'SimpleQueue',
+]
 
 
 try:
@@ -25,6 +33,10 @@ class Full(Exception):
     pass
 
 
+class ShutDown(Exception):
+    '''Raised when put/get with shut-down queue.'''
+
+
 class Queue:
     '''Create a queue object with a given maximum size.
 
@@ -54,6 +66,9 @@ def __init__(self, maxsize=0):
         self.all_tasks_done = threading.Condition(self.mutex)
         self.unfinished_tasks = 0
 
+        # Queue shutdown state
+        self.is_shutdown = False
+
     def task_done(self):
         '''Indicate that a formerly enqueued task is complete.
 
@@ -65,6 +80,9 @@ def task_done(self):
         have been processed (meaning that a task_done() call was received
         for every item that had been put() into the queue).
 
+        shutdown(immediate=True) calls task_done() for each remaining item in
+        the queue.
+
         Raises a ValueError if called more times than there were items
         placed in the queue.
         '''
@@ -129,15 +147,21 @@ def put(self, item, block=True, timeout=None):
         Otherwise ('block' is false), put an item on the queue if a free slot
         is immediately available, else raise the Full exception ('timeout'
         is ignored in that case).
+
+        Raises ShutDown if the queue has been shut down.
         '''
         with self.not_full:
+            if self.is_shutdown:
+                raise ShutDown
             if self.maxsize > 0:
                 if not block:
                     if self._qsize() >= self.maxsize:
                         raise Full
                 elif timeout is None:
                     while self._qsize() >= self.maxsize:
                         self.not_full.wait()
+                        if self.is_shutdown:
+                            raise ShutDown
                 elif timeout < 0:
                     raise ValueError("'timeout' must be a non-negative number")
                 else:
@@ -147,6 +171,8 @@ def put(self, item, block=True, timeout=None):
                         if remaining <= 0.0:
                             raise Full
                         self.not_full.wait(remaining)
+                        if self.is_shutdown:
+                            raise ShutDown
             self._put(item)
             self.unfinished_tasks += 1
             self.not_empty.notify()
@@ -161,14 +187,21 @@ def get(self, block=True, timeout=None):
         Otherwise ('block' is false), return an item if one is immediately
         available, else raise the Empty exception ('timeout' is ignored
         in that case).
+
+        Raises ShutDown if the queue has been shut down and is empty,
+        or if the queue has been shut down immediately.
         '''
         with self.not_empty:
+            if self.is_shutdown and not self._qsize():
+                raise ShutDown
             if not block:
                 if not self._qsize():
                     raise Empty
             elif timeout is None:
                 while not self._qsize():
                     self.not_empty.wait()
+                    if self.is_shutdown and not self._qsize():
+                        raise ShutDown
             elif timeout < 0:
                 raise ValueError("'timeout' must be a non-negative number")
             else:
@@ -178,6 +211,8 @@ def get(self, block=True, timeout=None):
                     if remaining <= 0.0:
                         raise Empty
                     self.not_empty.wait(remaining)
+                    if self.is_shutdown and not self._qsize():
+                        raise ShutDown
             item = self._get()
             self.not_full.notify()
             return item
@@ -198,6 +233,29 @@ def get_nowait(self):
         '''
         return self.get(block=False)
 
+    def shutdown(self, immediate=False):
+        '''Shut-down the queue, making queue gets and puts raise ShutDown.
+
+        By default, gets will only raise once the queue is empty. Set
+        'immediate' to True to make gets raise immediately instead.
+
+        All blocked callers of put() and get() will be unblocked. If
+        'immediate', a task is marked as done for each item remaining in
+        the queue, which may unblock callers of join().
+        '''
+        with self.mutex:
+            self.is_shutdown = True
+            if immediate:
+                while self._qsize():
+                    self._get()
+                    if self.unfinished_tasks > 0:
+                        self.unfinished_tasks -= 1
+                # release all blocked threads in `join()`
+                self.all_tasks_done.notify_all()
+            # All getters need to re-check queue-empty to raise ShutDown
+            self.not_empty.notify_all()
+            self.not_full.notify_all()
+
     # Override these methods to implement other queue organizations
     # (e.g. stack or priority queue).
     # These will only be called with appropriate locks held