From 806389117f96414a84678a72d6c8996337102485 Mon Sep 17 00:00:00 2001 From: Chris Lenk Date: Mon, 20 Jul 2020 00:24:36 -0400 Subject: [PATCH] Allow mixing single objects and lists in bundles ...in bundle constructor Related: #429. --- stix2/v20/bundle.py | 12 ++++++++---- stix2/v21/bundle.py | 12 ++++++++---- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/stix2/v20/bundle.py b/stix2/v20/bundle.py index 25b948d..6a663d6 100644 --- a/stix2/v20/bundle.py +++ b/stix2/v20/bundle.py @@ -26,10 +26,14 @@ class Bundle(_STIXBase20): def __init__(self, *args, **kwargs): # Add any positional arguments to the 'objects' kwarg. if args: - if isinstance(args[0], list): - kwargs['objects'] = args[0] + list(args[1:]) + kwargs.get('objects', []) - else: - kwargs['objects'] = list(args) + kwargs.get('objects', []) + obj_list = [] + for arg in args: + if isinstance(arg, list): + obj_list = obj_list + arg + else: + obj_list.append(arg) + + kwargs['objects'] = obj_list + kwargs.get('objects', []) self._allow_custom = kwargs.get('allow_custom', False) self._properties['objects'].contained.allow_custom = kwargs.get('allow_custom', False) diff --git a/stix2/v21/bundle.py b/stix2/v21/bundle.py index a1ca054..5497da5 100644 --- a/stix2/v21/bundle.py +++ b/stix2/v21/bundle.py @@ -23,10 +23,14 @@ class Bundle(_STIXBase21): def __init__(self, *args, **kwargs): # Add any positional arguments to the 'objects' kwarg. if args: - if isinstance(args[0], list): - kwargs['objects'] = args[0] + list(args[1:]) + kwargs.get('objects', []) - else: - kwargs['objects'] = list(args) + kwargs.get('objects', []) + obj_list = [] + for arg in args: + if isinstance(arg, list): + obj_list = obj_list + arg + else: + obj_list.append(arg) + + kwargs['objects'] = obj_list + kwargs.get('objects', []) self._allow_custom = kwargs.get('allow_custom', False) self._properties['objects'].contained.allow_custom = kwargs.get('allow_custom', False)