Allow shorter syntax for creating relationships.

stix2.1
Greg Back 2017-01-18 15:14:56 -08:00
parent fd548a5f41
commit 742d9645d6
2 changed files with 17 additions and 1 deletions

View File

@ -194,7 +194,8 @@ class Relationship(_STIXBase):
'target_ref',
]
def __init__(self, **kwargs):
def __init__(self, source_ref=None, relationship_type=None, target_ref=None,
**kwargs):
# TODO:
# - created_by_ref
# - revoked
@ -204,6 +205,13 @@ class Relationship(_STIXBase):
# - description
if source_ref and not kwargs.get('source_ref'):
kwargs['source_ref'] = source_ref
if relationship_type and not kwargs.get('relationship_type'):
kwargs['relationship_type'] = relationship_type
if target_ref and not kwargs.get('target_ref'):
kwargs['target_ref'] = target_ref
# TODO: do we care about the performance penalty of creating this
# if we won't need it?
now = datetime.now(tz=pytz.UTC)

View File

@ -336,3 +336,11 @@ def test_create_relationship_from_objects_rather_than_ids(indicator, malware):
assert relationship.relationship_type == 'indicates'
assert relationship.source_ref == INDICATOR_ID
assert relationship.target_ref == MALWARE_ID
def test_create_relationship_with_positional_args(indicator, malware):
relationship = stix2.Relationship(indicator, 'indicates', malware)
assert relationship.relationship_type == 'indicates'
assert relationship.source_ref == INDICATOR_ID
assert relationship.target_ref == MALWARE_ID