util: add a helper class to compute digests
authorMike Hommey <mh@glandium.org>
Thu, 16 Oct 2014 17:02:51 +0900
changeset 22962 4d58f4083148
parent 22961 a67ea4959ef5
child 22963 56e04741bbf1
util: add a helper class to compute digests It is going to be used for the remote-changegroup feature in bundle2.
mercurial/util.py
--- a/mercurial/util.py	Thu Oct 16 16:03:04 2014 +0900
+++ b/mercurial/util.py	Thu Oct 16 17:02:51 2014 +0900
@@ -117,6 +117,72 @@
     md5 = _md5
     return _md5(s)
 
+DIGESTS = {
+    'md5': md5,
+    'sha1': sha1,
+}
+# List of digest types from strongest to weakest
+DIGESTS_BY_STRENGTH = ['sha1', 'md5']
+
+try:
+    import hashlib
+    DIGESTS.update({
+        'sha512': hashlib.sha512,
+    })
+    DIGESTS_BY_STRENGTH.insert(0, 'sha512')
+except ImportError:
+    pass
+
+for k in DIGESTS_BY_STRENGTH:
+    assert k in DIGESTS
+
+class digester(object):
+    """helper to compute digests.
+
+    This helper can be used to compute one or more digests given their name.
+
+    >>> d = digester(['md5', 'sha1'])
+    >>> d.update('foo')
+    >>> [k for k in sorted(d)]
+    ['md5', 'sha1']
+    >>> d['md5']
+    'acbd18db4cc2f85cedef654fccc4a4d8'
+    >>> d['sha1']
+    '0beec7b5ea3f0fdbc95d0dd47f3c5bc275da8a33'
+    >>> digester.preferred(['md5', 'sha1'])
+    'sha1'
+    """
+
+    def __init__(self, digests, s=''):
+        self._hashes = {}
+        for k in digests:
+            if k not in DIGESTS:
+                raise Abort(_('unknown digest type: %s') % k)
+            self._hashes[k] = DIGESTS[k]()
+        if s:
+            self.update(s)
+
+    def update(self, data):
+        for h in self._hashes.values():
+            h.update(data)
+
+    def __getitem__(self, key):
+        if key not in DIGESTS:
+            raise Abort(_('unknown digest type: %s') % k)
+        return self._hashes[key].hexdigest()
+
+    def __iter__(self):
+        return iter(self._hashes)
+
+    @staticmethod
+    def preferred(supported):
+        """returns the strongest digest type in both supported and DIGESTS."""
+
+        for k in DIGESTS_BY_STRENGTH:
+            if k in supported:
+                return k
+        return None
+
 try:
     buffer = buffer
 except NameError: