Issue #18742: Rework the internal hashlib construtor to pave the road for ABCs.

2 files changed, 43 insertions, 33 deletions

diff --git a/Lib/hashlib.py b/Lib/hashlib.py
index 56a9360eac..77673f06eb 100644
--- a/Lib/hashlib.py
+++ b/Lib/hashlib.py
@@ -64,43 +64,42 @@ __all__ = __always_supported + ('new', 'algorithms_guaranteed',
                                 'algorithms_available', 'pbkdf2_hmac')
 
 
+__builtin_constructor_cache = {}
+
 def __get_builtin_constructor(name):
+    cache = __builtin_constructor_cache
+    constructor = cache.get(name)
+    if constructor is not None:
+        return constructor
     try:
         if name in ('SHA1', 'sha1'):
             import _sha1
-            return _sha1.sha1
+            cache['SHA1'] = cache['sha1'] = _sha1.sha1
         elif name in ('MD5', 'md5'):
             import _md5
-            return _md5.md5
+            cache['MD5'] = cache['md5'] = _md5.md5
         elif name in ('SHA256', 'sha256', 'SHA224', 'sha224'):
             import _sha256
-            bs = name[3:]
-            if bs == '256':
-                return _sha256.sha256
-            elif bs == '224':
-                return _sha256.sha224
+            cache['SHA224'] = cache['sha224'] = _sha256.sha224
+            cache['SHA256'] = cache['sha256'] = _sha256.sha256
         elif name in ('SHA512', 'sha512', 'SHA384', 'sha384'):
             import _sha512
-            bs = name[3:]
-            if bs == '512':
-                return _sha512.sha512
-            elif bs == '384':
-                return _sha512.sha384
+            cache['SHA384'] = cache['sha384'] = _sha512.sha384
+            cache['SHA512'] = cache['sha512'] = _sha512.sha512
         elif name in {'sha3_224', 'sha3_256', 'sha3_384', 'sha3_512',
                       'SHA3_224', 'SHA3_256', 'SHA3_384', 'SHA3_512'}:
             import _sha3
-            bs = name[5:]
-            if bs == '224':
-                return _sha3.sha3_224
-            elif bs == '256':
-                return _sha3.sha3_256
-            elif bs == '384':
-                return _sha3.sha3_384
-            elif bs == '512':
-                return _sha3.sha3_512
+            cache['SHA3_224'] = cache['sha3_224'] = _sha3.sha3_224
+            cache['SHA3_256'] = cache['sha3_256'] = _sha3.sha3_256
+            cache['SHA3_384'] = cache['sha3_384'] = _sha3.sha3_384
+            cache['SHA3_512'] = cache['sha3_512'] = _sha3.sha3_512
     except ImportError:
         pass  # no extension module, this hash is unsupported.
 
+    constructor = cache.get(name)
+    if constructor is not None:
+        return constructor
+
     raise ValueError('unsupported hash type ' + name)
 
 
diff --git a/Lib/test/test_hashlib.py b/Lib/test/test_hashlib.py
index 653f01c183..4a5ea7f1a9 100644
--- a/Lib/test/test_hashlib.py
+++ b/Lib/test/test_hashlib.py
@@ -84,26 +84,30 @@ class HashLibTestCase(unittest.TestCase):
                 if constructor:
                     constructors.add(constructor)
 
+        def add_builtin_constructor(name):
+            constructor = getattr(hashlib, "__get_builtin_constructor")(name)
+            self.constructors_to_test[name].add(constructor)
+
         _md5 = self._conditional_import_module('_md5')
         if _md5:
-            self.constructors_to_test['md5'].add(_md5.md5)
+            add_builtin_constructor('md5')
         _sha1 = self._conditional_import_module('_sha1')
         if _sha1:
-            self.constructors_to_test['sha1'].add(_sha1.sha1)
+            add_builtin_constructor('sha1')
         _sha256 = self._conditional_import_module('_sha256')
         if _sha256:
-            self.constructors_to_test['sha224'].add(_sha256.sha224)
-            self.constructors_to_test['sha256'].add(_sha256.sha256)
+            add_builtin_constructor('sha224')
+            add_builtin_constructor('sha256')
         _sha512 = self._conditional_import_module('_sha512')
         if _sha512:
-            self.constructors_to_test['sha384'].add(_sha512.sha384)
-            self.constructors_to_test['sha512'].add(_sha512.sha512)
+            add_builtin_constructor('sha384')
+            add_builtin_constructor('sha512')
         _sha3 = self._conditional_import_module('_sha3')
         if _sha3:
-            self.constructors_to_test['sha3_224'].add(_sha3.sha3_224)
-            self.constructors_to_test['sha3_256'].add(_sha3.sha3_256)
-            self.constructors_to_test['sha3_384'].add(_sha3.sha3_384)
-            self.constructors_to_test['sha3_512'].add(_sha3.sha3_512)
+            add_builtin_constructor('sha3_224')
+            add_builtin_constructor('sha3_256')
+            add_builtin_constructor('sha3_384')
+            add_builtin_constructor('sha3_512')
 
         super(HashLibTestCase, self).__init__(*args, **kwargs)
 
@@ -132,8 +136,10 @@ class HashLibTestCase(unittest.TestCase):
         self.assertRaises(TypeError, hashlib.new, 1)
 
     def test_get_builtin_constructor(self):
-        get_builtin_constructor = hashlib.__dict__[
-                '__get_builtin_constructor']
+        get_builtin_constructor = getattr(hashlib,
+                                          '__get_builtin_constructor')
+        builtin_constructor_cache = getattr(hashlib,
+                                            '__builtin_constructor_cache')
         self.assertRaises(ValueError, get_builtin_constructor, 'test')
         try:
             import _md5
@@ -141,6 +147,8 @@ class HashLibTestCase(unittest.TestCase):
             pass
         # This forces an ImportError for "import _md5" statements
         sys.modules['_md5'] = None
+        # clear the cache
+        builtin_constructor_cache.clear()
         try:
             self.assertRaises(ValueError, get_builtin_constructor, 'md5')
         finally:
@@ -149,6 +157,9 @@ class HashLibTestCase(unittest.TestCase):
             else:
                 del sys.modules['_md5']
         self.assertRaises(TypeError, get_builtin_constructor, 3)
+        constructor = get_builtin_constructor('md5')
+        self.assertIs(constructor, _md5.md5)
+        self.assertEqual(sorted(builtin_constructor_cache), ['MD5', 'md5'])
 
     def test_hexdigest(self):
         for cons in self.hash_constructors: