Refactor PrefixMap and PrefixList validation for clarity (#956)

* Refactor PrefixMap and PrefixList validation for clarity

* Add tests for validation of bad types

traits/tests/test_prefix_list.py

@@ -31,6 +31,18 @@ class A(HasTraits):
         with self.assertRaises(TraitError):
             a.foo = ''
 
+    def test_bad_types(self):
+        class A(HasTraits):
+            foo = PrefixList("zero", "one", "two", default_value="one")
+
+        a = A()
+
+        wrong_type = [[], (1, 2, 3), 1j, 2.3, 23, b"zero", None]
+        for value in wrong_type:
+            with self.subTest(value=value):
+                with self.assertRaises(TraitError):
+                    a.foo = value
+
     def test_repeated_prefix(self):
         class A(HasTraits):
             foo = PrefixList("abc1", "abc2")

traits/tests/test_prefix_map.py

@@ -18,11 +18,12 @@
 from traits.api import HasTraits, TraitError, PrefixMap, Undefined
 
 
+class Person(HasTraits):
+    married = PrefixMap({"yes": 1, "yeah": 1, "no": 0, "nah": 0})
+
+
 class TestPrefixMap(unittest.TestCase):
     def test_assignment(self):
-        class Person(HasTraits):
-            married = PrefixMap({"yes": 1, "yeah": 1, "no": 0, "nah": 0})
-
         person = Person()
 
         self.assertEqual(Undefined, person.married)
@@ -47,8 +48,14 @@ class Person(HasTraits):
         with self.assertRaises(TraitError):
             person.married = "ye"
 
-        with self.assertRaises(TraitError):
-            person.married = []
+    def test_bad_types(self):
+        person = Person()
+
+        wrong_type = [[], (1, 2, 3), 1j, 2.3, 23, b"not a string", None]
+        for value in wrong_type:
+            with self.subTest(value=value):
+                with self.assertRaises(TraitError):
+                    person.married = value
 
     def test_default(self):
         class Person(HasTraits):

traits/trait_types.py

@@ -2512,23 +2512,19 @@ def __init__(self, *values, **metadata):
         super().__init__(default, **metadata)
 
     def validate(self, object, name, value):
-        try:
-            if value not in self.values_:
-                match = None
-                n = len(value)
-                for key in self.values:
-                    if value == key[:n]:
-                        if match is not None:
-                            match = None
-                            break
-                        match = key
-                if match is None:
-                    self.error(object, name, value)
-                self.values_[value] = match
-            return self.values_[value]
-        except:
+        if not isinstance(value, str):
             self.error(object, name, value)
 
+        if value in self.values_:
+            return self.values_[value]
+
+        matches = [key for key in self.values if key.startswith(value)]
+        if len(matches) == 1:
+            self.values_[value] = match = matches[0]
+            return match
+
+        self.error(object, name, value)
+
     def info(self):
         return (
             " or ".join([repr(x) for x in self.values])
@@ -2901,24 +2897,18 @@ def __init__(self, map, **metadata):
         super().__init__(default_value, **metadata)
 
     def validate(self, object, name, value):
-        try:
-            if value in self._map:
-                return self._map[value]
-        except TypeError:
+        if not isinstance(value, str):
             self.error(object, name, value)
 
-        match = None
-        n = len(value)
-        for key in self.map.keys():
-            if value == key[:n]:
-                if match is not None:
-                    match = None
-                    break
-                match = key
-        if match is None:
-            self.error(object, name, value)
-        self._map[value] = match
-        return self._map[value]
+        if value in self._map:
+            return self._map[value]
+
+        matches = [key for key in self.map if key.startswith(value)]
+        if len(matches) == 1:
+            self._map[value] = match = matches[0]
+            return match
+
+        self.error(object, name, value)
 
     def mapped_value(self, value):
         """ Get the mapped value for a value. """