Allow adding validation to custom object types

_Custom.__init__() is a solution to recursion errors arising from using
a decorator. See https://stackoverflow.com/q/14739809/1013457
stix2.1
clenk 2017-06-13 10:26:43 -04:00
parent bcfb13f23c
commit d4e92dd813
2 changed files with 33 additions and 4 deletions

View File

@ -195,11 +195,35 @@ class Vulnerability(_STIXBase):
def CustomObject(type='x-custom-type', properties={}):
"""Custom STIX Object type decorator"""
"""Custom STIX Object type decorator
Example 1:
@CustomObject('type-name', {
'property1': StringProperty(required=True),
'property2': IntegerProperty(),
})
class MyNewObjectType():
pass
Supply an __init__() function to add any special validations to the custom
type. Don't call super().__init() though - doing so will cause an error.
Example 2:
@CustomObject('type-name', {
'property1': StringProperty(required=True),
'property2': IntegerProperty(),
})
class MyNewObjectType():
def __init__(self, property2=None, **kwargs):
if property2 and property2 < 10:
raise ValueError("'property2' is too small.")
"""
def custom_builder(cls):
class _Custom(_STIXBase):
class _Custom(cls, _STIXBase):
_type = type
_properties = COMMON_PROPERTIES.copy()
_properties.update({
@ -207,6 +231,11 @@ def CustomObject(type='x-custom-type', properties={}):
'type': TypeProperty(_type),
})
_properties.update(properties)
def __init__(self, **kwargs):
_STIXBase.__init__(self, **kwargs)
cls.__init__(self, **kwargs)
stix2._register_type(_Custom)
return _Custom

View File

@ -78,7 +78,7 @@ def test_parse_identity_custom_property(data):
})
class NewType():
def __init__(self, property2=None, **kwargs):
if property2 < 10:
if property2 and property2 < 10:
raise ValueError("'property2' is too small.")
@ -90,7 +90,7 @@ def test_custom_object_type():
NewType(property2=42)
with pytest.raises(ValueError):
NewType(property2=4)
NewType(property1='something', property2=4)
def test_parse_custom_object_type():