From 1ea9671b68e20bfe86739be72b8d8d990999d9a2 Mon Sep 17 00:00:00 2001 From: clenk Date: Wed, 5 Jul 2017 11:31:56 -0400 Subject: [PATCH] Allow passing a list to Bundle constructor Fix #26 --- stix2/bundle.py | 5 ++++- stix2/test/test_bundle.py | 24 ++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/stix2/bundle.py b/stix2/bundle.py index 85be3e1..b598ceb 100644 --- a/stix2/bundle.py +++ b/stix2/bundle.py @@ -17,6 +17,9 @@ class Bundle(_STIXBase): def __init__(self, *args, **kwargs): # Add any positional arguments to the 'objects' kwarg. if args: - kwargs['objects'] = kwargs.get('objects', []) + list(args) + if isinstance(args[0], list): + kwargs['objects'] = args[0] + list(args[1:]) + kwargs.get('objects', []) + else: + kwargs['objects'] = list(args) + kwargs.get('objects', []) super(Bundle, self).__init__(**kwargs) diff --git a/stix2/test/test_bundle.py b/stix2/test/test_bundle.py index d52a955..9c2dd19 100644 --- a/stix2/test/test_bundle.py +++ b/stix2/test/test_bundle.py @@ -92,3 +92,27 @@ def test_create_bundle_with_positional_args(indicator, malware, relationship): bundle = stix2.Bundle(indicator, malware, relationship) assert str(bundle) == EXPECTED_BUNDLE + + +def test_create_bundle_with_positional_listarg(indicator, malware, relationship): + bundle = stix2.Bundle([indicator, malware, relationship]) + + assert str(bundle) == EXPECTED_BUNDLE + + +def test_create_bundle_with_listarg_and_positional_arg(indicator, malware, relationship): + bundle = stix2.Bundle([indicator, malware], relationship) + + assert str(bundle) == EXPECTED_BUNDLE + + +def test_create_bundle_with_listarg_and_kwarg(indicator, malware, relationship): + bundle = stix2.Bundle([indicator, malware], objects=[relationship]) + + assert str(bundle) == EXPECTED_BUNDLE + + +def test_create_bundle_with_arg_listarg_and_kwarg(indicator, malware, relationship): + bundle = stix2.Bundle([indicator], malware, objects=[relationship]) + + assert str(bundle) == EXPECTED_BUNDLE