Issue #18882: Add threading.main_thread() function.

4 files changed, 102 insertions, 23 deletions

diff --git a/Doc/library/threading.rst b/Doc/library/threading.rst
index 5eb0069..9aec63d 100644
--- a/Doc/library/threading.rst
+++ b/Doc/library/threading.rst
@@ -57,6 +57,15 @@ This module defines the following functions:
    and threads that have not yet been started.
 
 
+.. function:: main_thread()
+
+   Return the main :class:`Thread` object.  In normal conditions, the
+   main thread is the thread from which the Python interpreter was
+   started.
+
+   .. versionadded:: 3.4
+
+
 .. function:: settrace(func)
 
    .. index:: single: trace function
diff --git a/Lib/test/test_threading.py b/Lib/test/test_threading.py
index 29ce039..971a635 100644
--- a/Lib/test/test_threading.py
+++ b/Lib/test/test_threading.py
@@ -21,6 +21,15 @@ import subprocess
 
 from test import lock_tests
 
+
+# Between fork() and exec(), only async-safe functions are allowed (issues
+# #12316 and #11870), and fork() from a worker thread is known to trigger
+# problems with some operating systems (issue #3863): skip problematic tests
+# on platforms known to behave badly.
+platforms_to_skip = ('freebsd4', 'freebsd5', 'freebsd6', 'netbsd5',
+                     'hp-ux11')
+
+
 # A trivial mutable counter.
 class Counter(object):
     def __init__(self):
@@ -468,15 +477,70 @@ class ThreadTests(BaseTestCase):
                 pid, status = os.waitpid(pid, 0)
                 self.assertEqual(0, status)
 
+    def test_main_thread(self):
+        main = threading.main_thread()
+        self.assertEqual(main.name, 'MainThread')
+        self.assertEqual(main.ident, threading.current_thread().ident)
+        self.assertEqual(main.ident, threading.get_ident())
 
-class ThreadJoinOnShutdown(BaseTestCase):
+        def f():
+            self.assertNotEqual(threading.main_thread().ident,
+                                threading.current_thread().ident)
+        th = threading.Thread(target=f)
+        th.start()
+        th.join()
+
+    @unittest.skipUnless(hasattr(os, 'fork'), "test needs os.fork()")
+    @unittest.skipUnless(hasattr(os, 'waitpid'), "test needs os.waitpid()")
+    def test_main_thread_after_fork(self):
+        code = """if 1:
+            import os, threading
+
+            pid = os.fork()
+            if pid == 0:
+                main = threading.main_thread()
+                print(main.name)
+                print(main.ident == threading.current_thread().ident)
+                print(main.ident == threading.get_ident())
+            else:
+                os.waitpid(pid, 0)
+        """
+        _, out, err = assert_python_ok("-c", code)
+        data = out.decode().replace('\r', '')
+        self.assertEqual(err, b"")
+        self.assertEqual(data, "MainThread\nTrue\nTrue\n")
+
+    @unittest.skipIf(sys.platform in platforms_to_skip, "due to known OS bug")
+    @unittest.skipUnless(hasattr(os, 'fork'), "test needs os.fork()")
+    @unittest.skipUnless(hasattr(os, 'waitpid'), "test needs os.waitpid()")
+    def test_main_thread_after_fork_from_nonmain_thread(self):
+        code = """if 1:
+            import os, threading, sys
+
+            def f():
+                pid = os.fork()
+                if pid == 0:
+                    main = threading.main_thread()
+                    print(main.name)
+                    print(main.ident == threading.current_thread().ident)
+                    print(main.ident == threading.get_ident())
+                    # stdout is fully buffered because not a tty,
+                    # we have to flush before exit.
+                    sys.stdout.flush()
+                else:
+                    os.waitpid(pid, 0)
 
-    # Between fork() and exec(), only async-safe functions are allowed (issues
-    # #12316 and #11870), and fork() from a worker thread is known to trigger
-    # problems with some operating systems (issue #3863): skip problematic tests
-    # on platforms known to behave badly.
-    platforms_to_skip = ('freebsd4', 'freebsd5', 'freebsd6', 'netbsd5',
-                         'hp-ux11')
+            th = threading.Thread(target=f)
+            th.start()
+            th.join()
+        """
+        _, out, err = assert_python_ok("-c", code)
+        data = out.decode().replace('\r', '')
+        self.assertEqual(err, b"")
+        self.assertEqual(data, "Thread-1\nTrue\nTrue\n")
+
+
+class ThreadJoinOnShutdown(BaseTestCase):
 
     def _run_and_join(self, script):
         script = """if 1:
diff --git a/Lib/threading.py b/Lib/threading.py
index 21bc781..2c06d99 100644
--- a/Lib/threading.py
+++ b/Lib/threading.py
@@ -840,20 +840,6 @@ class _MainThread(Thread):
         with _active_limbo_lock:
             _active[self._ident] = self
 
-    def _exitfunc(self):
-        self._stop()
-        t = _pickSomeNonDaemonThread()
-        while t:
-            t.join()
-            t = _pickSomeNonDaemonThread()
-        self._delete()
-
-def _pickSomeNonDaemonThread():
-    for t in enumerate():
-        if not t.daemon and t.is_alive():
-            return t
-    return None
-
 
 # Dummy thread class to represent threads not started here.
 # These aren't garbage collected when they die, nor can they be waited for.
@@ -915,7 +901,24 @@ from _thread import stack_size
 # and make it available for the interpreter
 # (Py_Main) as threading._shutdown.
 
-_shutdown = _MainThread()._exitfunc
+_main_thread = _MainThread()
+
+def _shutdown():
+    _main_thread._stop()
+    t = _pickSomeNonDaemonThread()
+    while t:
+        t.join()
+        t = _pickSomeNonDaemonThread()
+    _main_thread._delete()
+
+def _pickSomeNonDaemonThread():
+    for t in enumerate():
+        if not t.daemon and t.is_alive():
+            return t
+    return None
+
+def main_thread():
+    return _main_thread
 
 # get thread-local implementation, either from the thread
 # module, or from the python fallback
@@ -933,12 +936,13 @@ def _after_fork():
 
     # Reset _active_limbo_lock, in case we forked while the lock was held
     # by another (non-forked) thread.  http://bugs.python.org/issue874900
-    global _active_limbo_lock
+    global _active_limbo_lock, _main_thread
     _active_limbo_lock = _allocate_lock()
 
     # fork() only copied the current thread; clear references to others.
     new_active = {}
     current = current_thread()
+    _main_thread = current
     with _active_limbo_lock:
         for thread in _enumerate():
             # Any lock/condition variable may be currently locked or in an
diff --git a/Misc/NEWS b/Misc/NEWS
index 9539d62..5d4f75d 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -54,6 +54,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #18882: Add threading.main_thread() function.
+
 - Issue #18901: The sunau getparams method now returns a namedtuple rather than
   a plain tuple.  Patch by Claudiu Popa.