gh-128432: Add umask context manager to shutil

Doc/library/shutil.rst

@@ -877,6 +877,77 @@ 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, 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_context(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:: next
+
+In this example, we use a :class:`umask_context`, within which we create
+a file that only the user can access:
+
+.. code-block:: pycon
+
+   >>> from shutil import umask_context
+   >>> with umask_context(0o077):
+   ...     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 guest guest 20 Jan  1 23:45 my-secret-file
+
+Using :class:`umask_context` 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.
+
+It also allows you to write code that creates files, in a way that is agnostic
+of permissions -- that is without the need to pass a custom ``mode`` keyword
+argument to :func:`open` every time.
+
+In this example we create files with a function ``touch_file`` which uses the
+:func:`open` built-in without setting ``mode``. Permissions are managed by
+:class:`umask_context`:
+
+.. code-block:: pycon
+
+   >>> from shutil import umask_context
+   >>>
+   >>> def touch_file(path):
+   ...    with open(path, "a"):
+   ...        pass
+   ...
+   >>> touch_file("normal-file")
+   >>> with umask_context(0o077):
+   ...    touch_file("private-file")
+
+.. code-block:: shell-session
+
+   $ ls -l normal-file private-file
+   -rw-r--r-- 1 guest guest 0 Jan  1 23:45 normal-file
+   -rw------- 1 guest guest 0 Jan  1 23:45 private-file
+
 .. _`fcopyfile`:
    http://www.manpagez.com/man/3/copyfile/
 

Doc/whatsnew/3.14.rst

@@ -588,6 +588,14 @@ pydoc
   (Contributed by Jelle Zijlstra in :gh:`101552`.)
 
 
+shutil
+------
+
+* Add a new context manager :class:`~shutil.umask_context`, within which a
+  process's umask is temporarily changed
+  (Contributed by Jay Berry in :gh:`128433`.)
+
+
 ssl
 ---
 

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_context"]
            # 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_context(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_context)
 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_context']
         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 TestUmaskContext(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_context(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_context(target_mask_1), umask_context(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_context(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` context manager to :mod:`shutil`, to
+provide a stdlib context manager for changing a process's umask