bpo-31386: Custom wrap_bio and wrap_socket type (#3426)

SSLSocket.wrap_bio() and SSLSocket.wrap_socket() hard-code SSLObject and
SSLSocket as return types. In the light of future deprecation of
ssl.wrap_socket() module function and direct instantiation of SSLSocket,
it is desirable to make the return type of SSLSocket.wrap_bio() and
SSLSocket.wrap_socket() customizable.

Signed-off-by: Christian Heimes <christian@python.org>

1 files changed, 16 insertions, 0 deletions

diff --git a/Lib/test/test_ssl.py b/Lib/test/test_ssl.py
index 523322d..fb5958f 100644
--- a/Lib/test/test_ssl.py
+++ b/Lib/test/test_ssl.py
@@ -1359,6 +1359,22 @@ class ContextTests(unittest.TestCase):
         self.assertFalse(ctx.check_hostname)
         self.assertEqual(ctx.verify_mode, ssl.CERT_NONE)
 
+    def test_context_custom_class(self):
+        class MySSLSocket(ssl.SSLSocket):
+            pass
+
+        class MySSLObject(ssl.SSLObject):
+            pass
+
+        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        ctx.sslsocket_class = MySSLSocket
+        ctx.sslobject_class = MySSLObject
+
+        with ctx.wrap_socket(socket.socket(), server_side=True) as sock:
+            self.assertIsInstance(sock, MySSLSocket)
+        obj = ctx.wrap_bio(ssl.MemoryBIO(), ssl.MemoryBIO())
+        self.assertIsInstance(obj, MySSLObject)
+
 
 class SSLErrorTests(unittest.TestCase):