From 20be80c4ea8c193941a80dc2103814295ca11d71 Mon Sep 17 00:00:00 2001
From: John Vandenberg <jayvdb@gmail.com>
Date: Fri, 6 Sep 2019 12:11:04 +0700
Subject: [PATCH] DEV: Create ReplaceSysIoContextManager

Extracts the action of replacing sys.std* handles into a
separate context manager, so that alternative implementations
can re-use this component with other streams, or create wrappers
that do not use this functionality.
---
 src/stdio_mgr/stdio_mgr.py | 23 +++++------------------
 src/stdio_mgr/types.py     | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/src/stdio_mgr/stdio_mgr.py b/src/stdio_mgr/stdio_mgr.py
index 19e0b3c..d338178 100644
--- a/src/stdio_mgr/stdio_mgr.py
+++ b/src/stdio_mgr/stdio_mgr.py
@@ -26,7 +26,6 @@
 
 """
 
-import sys
 from contextlib import suppress
 from io import (
     BufferedRandom,
@@ -38,7 +37,7 @@
     TextIOWrapper,
 )
 
-from .types import _MultiCloseContextManager
+from .types import _MultiCloseContextManager, ReplaceSysIoContextManager
 
 
 class _PersistedBytesIO(BytesIO):
@@ -267,7 +266,7 @@ class SafeCloseTeeStdin(_SafeCloseIOBase, TeeStdin):
     """
 
 
-class StdioManager(_MultiCloseContextManager):
+class StdioManager(ReplaceSysIoContextManager, _MultiCloseContextManager):
     r"""Substitute temporary text buffers for `stdio` in a managed context.
 
     Context manager.
@@ -324,22 +323,10 @@ def __new__(cls, in_str="", close=True):
 
         return self
 
-    def __enter__(self):
-        """Enter context, replacing sys stdio objects with capturing streams."""
-        self._prior_streams = (sys.stdin, sys.stdout, sys.stderr)
-
-        super().__enter__()
-
-        (sys.stdin, sys.stdout, sys.stderr) = self
-
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        """Exit context, closing files and restoring state of sys module."""
-        (sys.stdin, sys.stdout, sys.stderr) = self._prior_streams
-
+    def close(self):
+        """Close files only if requested."""
         if self._close:
-            super().__exit__(exc_type, exc_value, traceback)
+            return super().close()
 
 
 stdio_mgr = StdioManager
diff --git a/src/stdio_mgr/types.py b/src/stdio_mgr/types.py
index 2d293a6..c53cd99 100644
--- a/src/stdio_mgr/types.py
+++ b/src/stdio_mgr/types.py
@@ -157,3 +157,23 @@ def __exit__(self, exc_type, exc_value, traceback):
         """Exit context, closing all members."""
         self.close()
         return super().__exit__(exc_type, exc_value, traceback)
+
+
+class ReplaceSysIoContextManager(StdioTuple):
+    """Replace sys stdio with members of the tuple."""
+
+    def __enter__(self):
+        """Enter context, replacing sys stdio objects with capturing streams."""
+        self._prior_streams = (sys.stdin, sys.stdout, sys.stderr)
+
+        super().__enter__()
+
+        (sys.stdin, sys.stdout, sys.stderr) = self
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Exit context, restoring state of sys module."""
+        (sys.stdin, sys.stdout, sys.stderr) = self._prior_streams
+
+        return super().__exit__(exc_type, exc_value, traceback)