Merge pull request #430 from oasis-open/bundle-objs-and-lists

Allow mixing single objects and lists in bundle constructor
oasis-open · Jul 22, 2020 · bb82bee · bb82bee
2 parents c497038 + 8063891
commit bb82bee
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git 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
@@ -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)