diff --git a/pytm/pytm.py b/pytm/pytm.py index a711687..6ee4543 100644 --- a/pytm/pytm.py +++ b/pytm/pytm.py @@ -181,9 +181,14 @@ def __set__(self, instance, value): class varClassification(var): def __set__(self, instance, value): - if not isinstance(value, Classification): + if isinstance(value, str): + Classification[value.lstrip("Classification.")] + + elif isinstance(value, Classification): + super().__set__(instance, value) + + elif not isinstance(value, Classification): raise ValueError("expecting a Classification, got a {}".format(type(value))) - super().__set__(instance, value) class varLifetime(var): diff --git a/tests/input.json b/tests/input.json index 0cdcff8..71344eb 100644 --- a/tests/input.json +++ b/tests/input.json @@ -5,7 +5,8 @@ "onDuplicates": "IGNORE", "boundaries": [ { - "name": "Internet" + "name": "Internet", + "maxClassification": "Classification.UNKNOWN" }, { "name": "Server/DB"