Fix bug when create dataframe from compound arrays and add test

G-Node · Mar 16, 2020 · e4b02fe · e4b02fe
1 parent 38a96ab
commit e4b02fe
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/nixio/block.py b/nixio/block.py
@@ -336,12 +336,11 @@ def create_data_frame(self, name="", type_="", col_dict=None,
                 if data is not None and type(data[0]) == np.void:
                     col_dtype = data[0].dtype
                     for i, dt in enumerate(col_dtype.fields.values()):
-                        if dt[0] == np.dtype(str):
-                            cn = list(col_dtype.fields.keys())
-                            raw_dt = col_dtype.fields.values()
-                            raw_dt = list(raw_dt)
-                            raw_dt_list = [ele[0] for ele in raw_dt]
-                            col_dict = OrderedDict(zip(cn, raw_dt_list))
+                        cn = list(col_dtype.fields.keys())
+                        raw_dt = col_dtype.fields.values()
+                        raw_dt = list(raw_dt)
+                        raw_dt_list = [ele[0] for ele in raw_dt]
+                    col_dict = OrderedDict(zip(cn, raw_dt_list))
                     if len(col_dtype.fields.values()) != len(col_dict):
                         raise exceptions.DuplicateColumnName
 
@@ -358,6 +357,8 @@ def create_data_frame(self, name="", type_="", col_dict=None,
                     if any(issubclass(dt, st) for st in string_types) \
                             or issubclass(dt, np.string_):
                         col_dict[nam] = util.vlen_str_dtype
+                if 'U' in str(dt) or dt == np.string_:
+                    col_dict[nam] = util.vlen_str_dtype
             dt_arr = list(col_dict.items())
             col_dtype = np.dtype(dt_arr)
 

diff --git a/nixio/test/test_data_frame.py b/nixio/test/test_data_frame.py
@@ -200,3 +200,9 @@ def test_data_type(self):
         assert self.df1.dtype[4] == np.int32
         assert self.df1.dtype[0] != self.df1.dtype[4]
         assert self.df1.dtype[2] == self.df1.dtype[3]
+
+    def test_creation_without_name(self):
+        b = np.array([("a",1,2.2), ("b",2,3.3), ("c",3,4.4)], dtype=[('name', 'U10'), ("id", 'i4'), ('val', 'f4')])
+        df = self.block.create_data_frame("without_name", "test", data=b)
+        assert list(df.column_names) == ["name", "id", "val"]
+        assert list(df["name"]) == ["a", "b", "c"]