Add new hashlib classes

astroid/brain/brain_hashlib.py

@@ -2,11 +2,13 @@
 # Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
 # For details: https://github.com/PyCQA/astroid/blob/master/COPYING.LESSER
+import sys
 import six
 import astroid
+PY36 = sys.version_info >= (3, 6)
 def _hashlib_transform():
     template = '''
@@ -29,7 +31,11 @@ def _hashlib_transform():
       def digest_size(self):
         return 1
     '''
-    algorithms = ('md5', 'sha1', 'sha224', 'sha256', 'sha384', 'sha512')
+    algorithms = ['md5', 'sha1', 'sha224', 'sha256', 'sha384', 'sha512']
+    if PY36:
+        algorithms += [
+            'sha3_224', 'sha3_256', 'sha3_384', 'sha3_512', 'shake_128', 'shake_256'
+        ]
     classes = "".join(
         template % {'name': hashfunc, 'digest': 'b""' if six.PY3 else '""'}
         for hashfunc in algorithms)
@@ -37,4 +43,3 @@ def _hashlib_transform():
 astroid.register_module_extender(astroid.MANAGER, 'hashlib', _hashlib_transform)
-

astroid/tests/unittest_brain.py

@@ -63,21 +63,31 @@ import astroid
 class HashlibTest(unittest.TestCase):
+    def _assert_hashlib_class(self, class_obj):
+        self.assertIn('update', class_obj)
+        self.assertIn('digest', class_obj)
+        self.assertIn('hexdigest', class_obj)
+        self.assertIn('block_size', class_obj)
+        self.assertIn('digest_size', class_obj)
+        self.assertEqual(len(class_obj['__init__'].args.args), 2)
+        self.assertEqual(len(class_obj['__init__'].args.defaults), 1)
+        self.assertEqual(len(class_obj['update'].args.args), 2)
+        self.assertEqual(len(class_obj['digest'].args.args), 1)
+        self.assertEqual(len(class_obj['hexdigest'].args.args), 1)
+
     def test_hashlib(self):
         """Tests that brain extensions for hashlib work."""
         hashlib_module = MANAGER.ast_from_module_name('hashlib')
         for class_name in ['md5', 'sha1']:
             class_obj = hashlib_module[class_name]
-            self.assertIn('update', class_obj)
-            self.assertIn('digest', class_obj)
-            self.assertIn('hexdigest', class_obj)
-            self.assertIn('block_size', class_obj)
-            self.assertIn('digest_size', class_obj)
-            self.assertEqual(len(class_obj['__init__'].args.args), 2)
-            self.assertEqual(len(class_obj['__init__'].args.defaults), 1)
-            self.assertEqual(len(class_obj['update'].args.args), 2)
-            self.assertEqual(len(class_obj['digest'].args.args), 1)
-            self.assertEqual(len(class_obj['hexdigest'].args.args), 1)
+            self._assert_hashlib_class(class_obj)
+
+    @test_utils.require_version(minver='3.6')
+    def test_hashlib_py36(self):
+        hashlib_module = MANAGER.ast_from_module_name('hashlib')
+        for class_name in ['sha3_224', 'sha3_512', 'shake_128']:
+            class_obj = hashlib_module[class_name]
+            self._assert_hashlib_class(class_obj)
 class CollectionsDequeTests(unittest.TestCase):