gh-128432: Add umask context manager to shutil

Doc/library/shutil.rst

@@ -877,6 +877,50 @@ Querying the size of the output terminal
       The ``fallback`` values are also used if :func:`os.get_terminal_size`
       returns zeroes.
 
+.. _shutil-context-managers:
+
+Context managers to modify process execution environments
+---------------------------------------------------------
+
+High-level :term:`context managers <context manager>` for changing a process's execution environment.
+
+.. warning::
+
+   These may change process-wide states, such as the current working directory,
+   and as such are not suitable for use in most threaded or async contexts.
+   They are also not suitable for most non-linear code execution, like generators,
+   where the program execution is temporarily relinquished. Unless explicitly
+   desired, you should not yield within these context managers.
+
+.. class:: umask(mask)
+
+   This is a simple wrapper around :func:`os.umask`. It changes the process's
+   umask upon entering and restores the old one on exit.
+
+   This context manager is :ref:`reentrant <reentrant-cms>`.
+
+   .. versionadded:: 3.14
+
+In this example, we use a :class:`umask` context manager, within which we
+create a file that only the user can access::
+
+    >>> from shutil import umask
+    >>> private_umask = umask(0o077)
+    >>> with private_umask:
+    ...     with open("my-secret-file", "w") as f:
+    ...         f.write("I ate all the cake!\n")
+
+The file's permissions are empty for anyone but the user:
+
+.. code-block:: shell-session
+
+    $ ls -l my-secret-file
+    -rw------- 1 jb2170 jb2170 20 Jan  2 01:23 my-secret-file
+
+Using :class:`umask` like this is better practice than first creating the file,
+and later changing its permissions with :func:`~os.chmod`, between which a
+period of time exists in which the file may have too lenient permissions.
+
 .. _`fcopyfile`:
    http://www.manpagez.com/man/3/copyfile/
 

Lib/shutil.py

@@ -10,6 +10,7 @@
 import fnmatch
 import collections
 import errno
+from contextlib import AbstractContextManager
 
 try:
     import zlib
@@ -61,7 +62,7 @@
            "get_unpack_formats", "register_unpack_format",
            "unregister_unpack_format", "unpack_archive",
            "ignore_patterns", "chown", "which", "get_terminal_size",
-           "SameFileError"]
+           "SameFileError", "umask"]
            # disk_usage is added later, if available on the platform
 
 class Error(OSError):
@@ -1581,6 +1582,21 @@ def which(cmd, mode=os.F_OK | os.X_OK, path=None):
                     return name
     return None
 
+
+class umask(AbstractContextManager):
+    """Non thread-safe context manager to change the process's umask."""
+
+    def __init__(self, mask):
+        self.mask = mask
+        self._old_mask = []
+
+    def __enter__(self):
+        self._old_mask.append(os.umask(self.mask))
+
+    def __exit__(self, *excinfo):
+        os.umask(self._old_mask.pop())
+
+
 def __getattr__(name):
     if name == "ExecError":
         import warnings

Lib/test/test_shutil.py

@@ -21,7 +21,7 @@
                     get_archive_formats, Error, unpack_archive,
                     register_unpack_format, RegistryError,
                     unregister_unpack_format, get_unpack_formats,
-                    SameFileError, _GiveupOnFastCopy)
+                    SameFileError, _GiveupOnFastCopy, umask)
 import tarfile
 import zipfile
 try:
@@ -3458,13 +3458,65 @@ def test_module_all_attribute(self):
                       'unregister_archive_format', 'get_unpack_formats',
                       'register_unpack_format', 'unregister_unpack_format',
                       'unpack_archive', 'ignore_patterns', 'chown', 'which',
-                      'get_terminal_size', 'SameFileError']
+                      'get_terminal_size', 'SameFileError', 'umask']
         if hasattr(os, 'statvfs') or os.name == 'nt':
             target_api.append('disk_usage')
         self.assertEqual(set(shutil.__all__), set(target_api))
         with self.assertWarns(DeprecationWarning):
             from shutil import ExecError
 
 
+@unittest.skipIf(os.name != "posix" or support.is_wasi or support.is_emscripten,
+                 "need proper os.umask()")
+class TestUmask(unittest.TestCase):
+    # make target masks in here sufficiently exotic, away from 0o022
+
+    mask_private = 0o777
+    mask_public  = 0o000
+
+    def get_mask(self):
+        os.umask(mask := os.umask(0))
+        return mask
+
+    def test_simple(self):
+        old_mask = self.get_mask()
+        target_mask = self.mask_private
+        self.assertNotEqual(old_mask, target_mask)
+
+        with umask(target_mask):
+            self.assertEqual(self.get_mask(), target_mask)
+        self.assertEqual(self.get_mask(), old_mask)
+
+    def test_reentrant(self):
+        old_mask = self.get_mask()
+        target_mask_1 = self.mask_private
+        target_mask_2 = self.mask_public
+        self.assertNotIn(old_mask, (target_mask_1, target_mask_2))
+        umask1, umask2 = umask(target_mask_1), umask(target_mask_2)
+
+        with umask1:
+            self.assertEqual(self.get_mask(), target_mask_1)
+            with umask2:
+                self.assertEqual(self.get_mask(), target_mask_2)
+                with umask1:
+                    self.assertEqual(self.get_mask(), target_mask_1)
+                self.assertEqual(self.get_mask(), target_mask_2)
+            self.assertEqual(self.get_mask(), target_mask_1)
+        self.assertEqual(self.get_mask(), old_mask)
+
+    def test_exception(self):
+        old_mask = self.get_mask()
+        target_mask = self.mask_private
+        self.assertNotEqual(old_mask, target_mask)
+
+        try:
+            with umask(target_mask):
+                self.assertEqual(self.get_mask(), target_mask)
+                raise RuntimeError("boom")
+        except RuntimeError as re:
+            self.assertEqual(str(re), "boom")
+        self.assertEqual(self.get_mask(), old_mask)
+
+
 if __name__ == '__main__':
     unittest.main()

Misc/NEWS.d/next/Library/2025-01-03-01-12-47.gh-issue-128432.8kg5DK.rst

@@ -0,0 +1,2 @@
+Add a :class:`~shutil.umask` context manager to :mod:`shutil`, to provide a
+stdlib context manager for changing a process's umask