Merge pull request #417 from chisholm/fix_list_property

Fix ListProperty

stix2/properties.py

@@ -189,15 +189,29 @@ class ListProperty(Property):
 
     def __init__(self, contained, **kwargs):
         """
-        ``contained`` should be a function which returns an object from the value.
+        ``contained`` should be a Property class or instance, or a _STIXBase
+        subclass.
         """
-        if inspect.isclass(contained) and issubclass(contained, Property):
-            # If it's a class and not an instance, instantiate it so that
-            # clean() can be called on it, and ListProperty.clean() will
-            # use __call__ when it appends the item.
-            self.contained = contained()
-        else:
+        self.contained = None
+
+        if inspect.isclass(contained):
+            # Property classes are instantiated; _STIXBase subclasses are left
+            # as-is.
+            if issubclass(contained, Property):
+                self.contained = contained()
+            elif issubclass(contained, _STIXBase):
+                self.contained = contained
+
+        elif isinstance(contained, Property):
             self.contained = contained
+
+        if not self.contained:
+            raise TypeError(
+                "Invalid list element type: {}".format(
+                    str(contained),
+                ),
+            )
+
         super(ListProperty, self).__init__(**kwargs)
 
     def clean(self, value):
@@ -209,40 +223,28 @@ class ListProperty(Property):
         if isinstance(value, (_STIXBase, string_types)):
             value = [value]
 
-        result = []
-        for item in value:
-            try:
-                valid = self.contained.clean(item)
-            except ValueError:
-                raise
-            except AttributeError:
-                # type of list has no clean() function (eg. built in Python types)
-                # TODO Should we raise an error here?
-                valid = item
+        if isinstance(self.contained, Property):
+            result = [
+                self.contained.clean(item)
+                for item in value
+            ]
 
-            if type(self.contained) is EmbeddedObjectProperty:
-                obj_type = self.contained.type
-            elif type(self.contained).__name__ == "STIXObjectProperty":
-                # ^ this way of checking doesn't require a circular import
-                # valid is already an instance of a python-stix2 class; no need
-                # to turn it into a dictionary and then pass it to the class
-                # constructor again
-                result.append(valid)
-                continue
-            elif type(self.contained) is DictionaryProperty:
-                obj_type = dict
-            else:
-                obj_type = self.contained
+        else:  # self.contained must be a _STIXBase subclass
+            result = []
+            for item in value:
+                if isinstance(item, self.contained):
+                    valid = item
+
+                elif isinstance(item, Mapping):
+                    # attempt a mapping-like usage...
+                    valid = self.contained(**item)
 
-            if isinstance(valid, Mapping):
-                try:
-                    valid._allow_custom
-                except AttributeError:
-                    result.append(obj_type(**valid))
                 else:
-                    result.append(obj_type(allow_custom=True, **valid))
-            else:
-                result.append(obj_type(valid))
+                    raise ValueError("Can't create a {} out of {}".format(
+                        self.contained._type, str(item),
+                    ))
+
+                result.append(valid)
 
         # STIX spec forbids empty lists
         if len(result) < 1:

stix2/test/v20/test_properties.py

@@ -3,8 +3,10 @@ import uuid
 import pytest
 
 import stix2
+import stix2.base
 from stix2.exceptions import (
     AtLeastOnePropertyError, CustomContentError, DictionaryKeyError,
+    ExtraPropertiesError,
 )
 from stix2.properties import (
     BinaryProperty, BooleanProperty, DictionaryProperty,
@@ -66,7 +68,7 @@ def test_fixed_property():
     assert p.clean(p.default())
 
 
-def test_list_property():
+def test_list_property_property_type():
     p = ListProperty(StringProperty)
 
     assert p.clean(['abc', 'xyz'])
@@ -74,6 +76,88 @@ def test_list_property():
         p.clean([])
 
 
+def test_list_property_property_type_custom():
+    class TestObj(stix2.base._STIXBase):
+        _type = "test"
+        _properties = {
+            "foo": StringProperty(),
+        }
+    p = ListProperty(EmbeddedObjectProperty(type=TestObj))
+
+    objs_custom = [
+        TestObj(foo="abc", bar=123, allow_custom=True),
+        TestObj(foo="xyz"),
+    ]
+
+    assert p.clean(objs_custom)
+
+    dicts_custom = [
+        {"foo": "abc", "bar": 123},
+        {"foo": "xyz"},
+    ]
+
+    # no opportunity to set allow_custom=True when using dicts
+    with pytest.raises(ExtraPropertiesError):
+        p.clean(dicts_custom)
+
+
+def test_list_property_object_type():
+    class TestObj(stix2.base._STIXBase):
+        _type = "test"
+        _properties = {
+            "foo": StringProperty(),
+        }
+    p = ListProperty(TestObj)
+
+    objs = [TestObj(foo="abc"), TestObj(foo="xyz")]
+    assert p.clean(objs)
+
+    dicts = [{"foo": "abc"}, {"foo": "xyz"}]
+    assert p.clean(dicts)
+
+
+def test_list_property_object_type_custom():
+    class TestObj(stix2.base._STIXBase):
+        _type = "test"
+        _properties = {
+            "foo": StringProperty(),
+        }
+    p = ListProperty(TestObj)
+
+    objs_custom = [
+        TestObj(foo="abc", bar=123, allow_custom=True),
+        TestObj(foo="xyz"),
+    ]
+
+    assert p.clean(objs_custom)
+
+    dicts_custom = [
+        {"foo": "abc", "bar": 123},
+        {"foo": "xyz"},
+    ]
+
+    # no opportunity to set allow_custom=True when using dicts
+    with pytest.raises(ExtraPropertiesError):
+        p.clean(dicts_custom)
+
+
+def test_list_property_bad_element_type():
+    with pytest.raises(TypeError):
+        ListProperty(1)
+
+
+def test_list_property_bad_value_type():
+    class TestObj(stix2.base._STIXBase):
+        _type = "test"
+        _properties = {
+            "foo": StringProperty(),
+        }
+
+    list_prop = ListProperty(TestObj)
+    with pytest.raises(ValueError):
+        list_prop.clean([1])
+
+
 def test_string_property():
     prop = StringProperty()