PR feedback: swap to dict with constants, better subtest parameterization, lint

Author: csm10495

Lib/concurrent/futures/process.py

@@ -45,7 +45,6 @@
 
 __author__ = 'Brian Quinlan ([email protected])'
 
-import enum
 import os
 from concurrent.futures import _base
 import queue
@@ -627,17 +626,13 @@ class BrokenProcessPool(_base.BrokenExecutor):
     while a future was in the running state.
     """
 
-class _TerminateOrKillOperation(enum.Enum):
-    """Enum for _terminate_or_kill_workers().
+_TERMINATE = "terminate"
+_KILL = "kill"
 
-    Used to determine the operation used by the
-    _terminate_or_kill_workers() method.
-    """
-    # Delegate to call process.terminate()
-    TERMINATE = 1
-
-    # Delegate to call process.kill()
-    KILL = 2
+_TERMINATE_OR_KILL_OPERATION = {
+    _TERMINATE,
+    _KILL
+}
 
 
 class ProcessPoolExecutor(_base.Executor):
@@ -869,24 +864,25 @@ def shutdown(self, wait=True, *, cancel_futures=False):
 
     shutdown.__doc__ = _base.Executor.shutdown.__doc__
 
-    def _terminate_or_kill_workers(self, operation: _TerminateOrKillOperation):
-        """Attempts to terminate or kill the executor's workers based off the given
-        operation. Iterates through all of the current processes and performs the
-        relevant task if the process is still alive.
+    def _terminate_or_kill_workers(self, operation):
+        """Attempts to terminate or kill the executor's workers based off the
+        given operation. Iterates through all of the current processes and
+        performs the relevant task if the process is still alive.
 
         After terminating workers, the pool will be in a broken state
         and no longer usable (for instance, new tasks should not be
         submitted).
         """
-        if operation not in _TerminateOrKillOperation:
+        if operation not in _TERMINATE_OR_KILL_OPERATION:
             raise ValueError(f"Unsupported operation: {operation}")
 
         processes = {}
         if self._processes:
             processes = self._processes.copy()
 
-        # shutdown will invalidate ._processes, so we copy it right before calling.
-        # If we waited here, we would deadlock if a process decides not to exit.
+        # shutdown will invalidate ._processes, so we copy it right before
+        # calling. If we waited here, we would deadlock if a process decides not
+        # to exit.
         self.shutdown(wait=False, cancel_futures=True)
 
         if not processes:
@@ -901,9 +897,9 @@ def _terminate_or_kill_workers(self, operation: _TerminateOrKillOperation):
                 continue
 
             try:
-                if operation == _TerminateOrKillOperation.TERMINATE:
+                if operation == _TERMINATE:
                     proc.terminate()
-                elif operation == _TerminateOrKillOperation.KILL:
+                elif operation == _KILL:
                     proc.kill()
             except ProcessLookupError:
                 # The process just ended before our signal
@@ -918,7 +914,7 @@ def terminate_workers(self):
         and no longer usable (for instance, new tasks should not be
         submitted).
         """
-        return self._terminate_or_kill_workers(operation=_TerminateOrKillOperation.TERMINATE)
+        return self._terminate_or_kill_workers(operation=_TERMINATE)
 
     def kill_workers(self):
         """Attempts to kill the executor's workers.
@@ -929,4 +925,4 @@ def kill_workers(self):
         and no longer usable (for instance, new tasks should not be
         submitted).
         """
-        return self._terminate_or_kill_workers(operation=_TerminateOrKillOperation.KILL)
+        return self._terminate_or_kill_workers(operation=_KILL)

Lib/test/test_concurrent_futures/test_process_pool.py

@@ -12,6 +12,7 @@
 
 from test import support
 from test.support import hashlib_helper
+from test.test_importlib.metadata.fixtures import parameterize
 
 from .executor import ExecutorTest, mul
 from .util import (
@@ -26,6 +27,11 @@ def __init__(self, mgr):
     def __del__(self):
         self.event.set()
 
+TERMINATE_OR_KILL_PARAMS = [
+    dict(function_name='terminate_workers'),
+    dict(function_name='kill_workers'),
+]
+
 def _put_sleep_put(queue):
     """ Used as part of test_process_pool_executor_terminate_workers """
     queue.put('started')
@@ -228,80 +234,80 @@ def mock_start_new_thread(func, *args, **kwargs):
                     list(executor.map(mul, [(2, 3)] * 10))
             executor.shutdown()
 
-    def test_process_pool_executor_terminate_kill_workers(self):
-        for function_name in ('terminate_workers', 'kill_workers'):
-            with self.subTest(function_name):
-                manager = multiprocessing.Manager()
-                q = manager.Queue()
-
-                with futures.ProcessPoolExecutor(max_workers=1) as executor:
-                    executor.submit(_put_sleep_put, q)
-
-                    # We should get started, but not finished since we'll terminate the workers just after
-                    self.assertEqual(q.get(timeout=5), 'started')
+    def test_process_pool_executor_terminate_workers(self):
+        with futures.ProcessPoolExecutor(max_workers=1) as executor:
+            executor._terminate_or_kill_workers = unittest.mock.Mock()
+            executor.terminate_workers()
 
-                    worker_process = list(executor._processes.values())[0]
-                    getattr(executor, function_name)()
-                    worker_process.join()
+        executor._terminate_or_kill_workers.assert_called_once_with(operation=futures.process._TERMINATE)
 
-                    if function_name == 'terminate_workers' or sys.platform == 'win32':
-                        # On windows, kill and terminate both send SIGTERM
-                        self.assertEqual(worker_process.exitcode, -signal.SIGTERM)
-                    elif function_name == 'kill_workers':
-                        self.assertEqual(worker_process.exitcode, -signal.SIGKILL)
+    def test_process_pool_executor_kill_workers(self):
+        with futures.ProcessPoolExecutor(max_workers=1) as executor:
+            executor._terminate_or_kill_workers = unittest.mock.Mock()
+            executor.kill_workers()
 
-                    self.assertRaises(queue.Empty, q.get, timeout=1)
+        executor._terminate_or_kill_workers.assert_called_once_with(operation=futures.process._KILL)
 
-    def test_process_pool_executor_terminate_kill_workers_dead_workers(self):
-        for function_name in ('terminate_workers', 'kill_workers'):
-            with self.subTest(function_name=function_name):
-                with futures.ProcessPoolExecutor(max_workers=1) as executor:
-                    future = executor.submit(os._exit, 1)
-                    self.assertRaises(BrokenProcessPool, future.result)
+    def test_process_pool_executor_terminate_or_kill_workers_invalid_op(self):
+        with futures.ProcessPoolExecutor(max_workers=1) as executor:
+            self.assertRaises(ValueError,
+                              executor._terminate_or_kill_workers,
+                              operation='invalid operation'),
 
-                    # even though the pool is broken, this shouldn't raise
-                    getattr(executor, function_name)()
+    @parameterize(*TERMINATE_OR_KILL_PARAMS)
+    def test_process_pool_executor_terminate_kill_workers(self, function_name):
+        manager = multiprocessing.Manager()
+        q = manager.Queue()
 
-    def test_process_pool_executor_terminate_kill_workers_not_started_yet(self):
-        for function_name in ('terminate_workers', 'kill_workers'):
-            with self.subTest(function_name=function_name):
-                context_with_mocked_process = multiprocessing.get_context()
-                with unittest.mock.patch.object(context_with_mocked_process, 'Process') as mock_process:
-                    with futures.ProcessPoolExecutor(max_workers=1, mp_context=context_with_mocked_process) as executor:
-                        # The worker has not been started yet, terminate/kill_workers should basically no-op
-                        getattr(executor, function_name)()
+        with futures.ProcessPoolExecutor(max_workers=1) as executor:
+            executor.submit(_put_sleep_put, q)
 
-                    mock_process.return_value.kill.assert_not_called()
-                    mock_process.return_value.terminate.assert_not_called()
+            # We should get started, but not finished since we'll terminate the
+            # workers just after
+            self.assertEqual(q.get(timeout=5), 'started')
 
-    def test_process_pool_executor_terminate_kill_workers_stops_pool(self):
-        for function_name in ('terminate_workers', 'kill_workers'):
-            with self.subTest(function_name=function_name):
-                with futures.ProcessPoolExecutor(max_workers=1) as executor:
-                    executor.submit(time.sleep, 0).result()
+            worker_process = list(executor._processes.values())[0]
+            getattr(executor, function_name)()
+            worker_process.join()
 
-                    getattr(executor, function_name)()
+            if function_name == 'terminate_workers' or \
+                sys.platform == 'win32':
+                # On windows, kill and terminate both send SIGTERM
+                self.assertEqual(worker_process.exitcode, -signal.SIGTERM)
+            elif function_name == 'kill_workers':
+                self.assertEqual(worker_process.exitcode, -signal.SIGKILL)
 
-                    self.assertRaises(RuntimeError, executor.submit, time.sleep, 0)
+            self.assertRaises(queue.Empty, q.get, timeout=1)
 
-    def test_process_pool_executor_terminate_workers(self):
+    @parameterize(*TERMINATE_OR_KILL_PARAMS)
+    def test_process_pool_executor_terminate_kill_workers_dead_workers(self, function_name):
         with futures.ProcessPoolExecutor(max_workers=1) as executor:
-            executor._terminate_or_kill_workers = unittest.mock.Mock()
-            executor.terminate_workers()
-
-        executor._terminate_or_kill_workers.assert_called_once_with(operation=futures.process._TerminateOrKillOperation.TERMINATE)
-
-    def test_process_pool_executor_kill_workers(self):
+            future = executor.submit(os._exit, 1)
+            self.assertRaises(BrokenProcessPool, future.result)
+
+            # even though the pool is broken, this shouldn't raise
+            getattr(executor, function_name)()
+
+    @parameterize(*TERMINATE_OR_KILL_PARAMS)
+    def test_process_pool_executor_terminate_kill_workers_not_started_yet(self, function_name):
+        ctx = self.get_context()
+        with unittest.mock.patch.object(ctx, 'Process') as mock_process:
+            with futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
+                # The worker has not been started yet, terminate/kill_workers
+                # should basically no-op
+                getattr(executor, function_name)()
+
+            mock_process.return_value.kill.assert_not_called()
+            mock_process.return_value.terminate.assert_not_called()
+
+    @parameterize(*TERMINATE_OR_KILL_PARAMS)
+    def test_process_pool_executor_terminate_kill_workers_stops_pool(self, function_name):
         with futures.ProcessPoolExecutor(max_workers=1) as executor:
-            executor._terminate_or_kill_workers = unittest.mock.Mock()
-            executor.kill_workers()
+            executor.submit(time.sleep, 0).result()
 
-        executor._terminate_or_kill_workers.assert_called_once_with(operation=futures.process._TerminateOrKillOperation.KILL)
-
-    def test_process_pool_executor_terminate_or_kill_workers_invalid_operation(self):
-        with futures.ProcessPoolExecutor(max_workers=1) as executor:
-            self.assertRaises(ValueError, executor._terminate_or_kill_workers, operation='invalid operation'),
+            getattr(executor, function_name)()
 
+            self.assertRaises(RuntimeError, executor.submit, time.sleep, 0)
 
 
 create_executor_tests(globals(), ProcessPoolExecutorTest,