Fix mapping of multidimensional columns

Author: smathot

datamatrix/_datamatrix/_multidimensionalcolumn.py

@@ -16,12 +16,13 @@
 You should have received a copy of the GNU General Public License
 along with datamatrix.  If not, see <http://www.gnu.org/licenses/>.
 """
+from datamatrix.py3compat import *
 import logging
-import os
 import weakref
-from datamatrix.py3compat import *
 from datamatrix import cfg
-from datamatrix._datamatrix._numericcolumn import NumericColumn, FloatColumn
+from datamatrix._datamatrix._mixedcolumn import MixedColumn
+from datamatrix._datamatrix._numericcolumn import NumericColumn, FloatColumn, \
+    IntColumn
 from datamatrix._datamatrix._datamatrix import DataMatrix
 from collections.abc import Sequence, Collection
 from collections import OrderedDict
@@ -348,11 +349,19 @@ def _map(self, fnc):
 
         # For a MultiDimensionalColumn, we need to make a special case, because
         # the shape of the new MultiDimensionalColumn may be different from
-        # the shape of the original column.
+        # the shape of the original column. The new column may even be a 
+        # different kind of column altogether.
         for i, cell in enumerate(self):
             a = fnc(cell)
             if not i:
-                newcol = self.__class__(self.dm, shape=len(a))
+                if isinstance(a, float):
+                    newcol = FloatColumn(self.dm)
+                elif isinstance(a, int):
+                    newcol = IntColumn(self.dm)
+                elif isinstance(a, str):
+                    newcol = MixedColumn(self.dm)
+                else:
+                    newcol = self.__class__(self.dm, shape=len(a))
             newcol[i] = a
         return newcol
 

testcases/test_functional.py

@@ -18,25 +18,42 @@
 """
 
 from datamatrix.py3compat import *
-from datamatrix import DataMatrix, MixedColumn, IntColumn, FloatColumn
+from datamatrix import DataMatrix, MixedColumn, IntColumn, FloatColumn, \
+    MultiDimensionalColumn
 from datamatrix import functional as fnc
-from testcases.test_tools import capture_stdout
+from testcases.test_tools import capture_stdout, check_series
 
 
-def test_map_():
+def test_map():
 
     for coltype in (MixedColumn, FloatColumn, IntColumn):
         dm = DataMatrix(length=2, default_col_type=coltype)
         dm.a = 1, 2
         dm.a = fnc.map_(lambda x: x*2, dm.a)
         assert dm.a == [2, 4]
         assert isinstance(dm.a, coltype)
+        dm.a = 1, 2
+        dm.a = dm.a @ (lambda x: x*2)
+        assert dm.a == [2, 4]
+        assert isinstance(dm.a, coltype)
         dm = fnc.map_(lambda **d: {'a' : 0}, dm)
         assert dm.a == [0, 0]
         assert isinstance(dm.a, coltype)
 
 
-def test_filter_():
+def test_map_multidimensional():
+    dm = DataMatrix(length=2)
+    dm.m = MultiDimensionalColumn(shape=(3,))
+    dm.m = [[1,2,3], [4,5,6]]
+    dm.mean = dm.m @ (lambda a: a.mean())
+    assert dm.mean == [2, 5]
+    dm.half = dm.m @ (lambda a: a / 2)
+    check_series(dm.half, [[0.5, 1., 1.5], [2, 2.5, 3.]])
+    dm.short = dm.m @ (lambda a: a[:2])
+    check_series(dm.short, [[1, 2], [4, 5]])
+
+
+def test_filter():
 
     dm = DataMatrix(length=4)
     dm.a = range(4)