diff --git a/Lib/hashlib.py b/Lib/hashlib.py index 56a9360eac6..77673f06eb2 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 653f01c183c..4a5ea7f1a94 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: