diff --git a/swh/model/model.py b/swh/model/model.py index b741ae7c9232db2d75a7a7aa7fec6935b07671ca..36f8ef3f1838fcf0ea8c632962cde906ee7db5bb 100644 --- a/swh/model/model.py +++ b/swh/model/model.py @@ -11,6 +11,7 @@ import attr import dateutil.parser from .identifiers import normalize_timestamp +from .hashutil import DEFAULT_ALGORITHMS SHA1_SIZE = 20 @@ -362,3 +363,8 @@ class Content(BaseModel): if content[field] is None: del content[field] return content + + def get_hash(self, hash_name): + if hash_name not in DEFAULT_ALGORITHMS: + raise ValueError('{} is not a valid hash name.'.format(hash_name)) + return getattr(self, hash_name) diff --git a/swh/model/tests/test_model.py b/swh/model/tests/test_model.py index 220ba324798e9a5f522483f731502f16e9950793..f91868577235614cd5733957a0af6d19a289c046 100644 --- a/swh/model/tests/test_model.py +++ b/swh/model/tests/test_model.py @@ -7,6 +7,7 @@ import copy from hypothesis import given +from swh.model.model import Content from swh.model.hypothesis_strategies import objects @@ -24,3 +25,11 @@ def test_todict_inverse_fromdict(objtype_and_obj): # Check the composition of from_dict and to_dict is the identity assert obj_as_dict == type(obj).from_dict(obj_as_dict).to_dict() + + +def test_content_get_hash(): + hashes = dict( + sha1=b'foo', sha1_git=b'bar', sha256=b'baz', blake2s256=b'qux') + c = Content(length=42, status='visible', **hashes) + for (hash_name, hash_) in hashes.items(): + assert c.get_hash(hash_name) == hash_