Use strings directly for casing as np.object w/o redundant StringHolder. (#1736)

yuslepukhin · houseroad · commit beefa15c3c20 · 2019-01-14T17:03:20.000-08:00
diff --git a/onnx/mapping.py b/onnx/mapping.py
@@ -7,20 +7,6 @@
 from typing import Text, Any
 import numpy as np  # type: ignore
 
-
-class StringHolder:
-    def __init__(self, s):    # type: (Text) -> None
-        self.text = s
-
-    def __eq__(self, other):  # type: (object) -> Any
-        if isinstance(other, self.__class__):
-            return self.text == getattr(other, 'text')
-        return False
-
-    def __ne__(self, other):  # type: (object) -> Any
-        return not self.__eq__(other)
-
-
 TENSOR_TYPE_TO_NP_TYPE = {
     int(TensorProto.FLOAT): np.dtype('float32'),
     int(TensorProto.UINT8): np.dtype('uint8'),
diff --git a/onnx/numpy_helper.py b/onnx/numpy_helper.py
@@ -8,7 +8,7 @@
 import numpy as np  # type: ignore
 from onnx import TensorProto
 from onnx import mapping
-from onnx.mapping import StringHolder
+from six import text_type
 from typing import Sequence, Any, Optional, Text, List
 
 
@@ -45,7 +45,7 @@ def to_array(tensor):  # type: (TensorProto) -> np.ndarray[Any]
 
     if tensor.data_type == TensorProto.STRING:
         utf8_strings = getattr(tensor, storage_field)
-        ss = list(StringHolder(s.decode('utf-8')) for s in utf8_strings)
+        ss = list(s.decode('utf-8') for s in utf8_strings)
         return np.asarray(ss).astype(np_dtype).reshape(dims)
 
     if tensor.HasField("raw_data"):
@@ -85,10 +85,10 @@ def from_array(arr, name=None):  # type: (np.ndarray[Any], Optional[Text]) -> Te
         # Special care for strings.
         tensor.data_type = mapping.NP_TYPE_TO_TENSOR_TYPE[arr.dtype]
         for e in arr:
-            if isinstance(e, StringHolder):
-                tensor.string_data.append(e.text.encode('utf-8'))
+            if isinstance(e, text_type):
+                tensor.string_data.append(e.encode('utf-8'))
             else:
-                raise NotImplementedError("Unrecognized object in the object array, expect StringHolder")
+                raise NotImplementedError("Unrecognized object in the object array, expect a string")
 
         return tensor
 
diff --git a/onnx/test/numpy_helper_test.py b/onnx/test/numpy_helper_test.py
@@ -6,7 +6,6 @@
 import numpy as np  # type: ignore
 
 from onnx import numpy_helper
-from onnx.mapping import StringHolder
 
 import unittest
 
@@ -53,8 +52,7 @@ def test_int64(self):  # type: () -> None
         self._test_numpy_helper_int_type(np.int64)
 
     def test_string(self):  # type: () -> None
-        strholder_list = list(StringHolder(s) for s in ['Amy', 'Billy', 'Cindy', 'David'])
-        a = np.array(strholder_list).astype(np.object)
+        a = np.array(['Amy', 'Billy', 'Cindy', 'David']).astype(np.object)
         tensor_def = numpy_helper.from_array(a, "test")
         self.assertEqual(tensor_def.name, "test")
         a_recover = numpy_helper.to_array(tensor_def)
