code cleanup, additional tests

Author: flexatone

src/tri_map.c

@@ -1135,7 +1135,6 @@ TriMap_map_dst_no_fill(TriMapObject *self, PyObject *arg) {
     return AK_TM_map_no_fill(self, from_src, array_from);
 }
 
-
 static inline PyObject *
 TriMap_map_merge(TriMapObject *tm, PyObject *args)
 {
@@ -1178,30 +1177,15 @@ TriMap_map_merge(TriMapObject *tm, PyObject *args)
         Py_DECREF(dtype); // not needed
         // will initialize to NULL, not None
         array_to = (PyArrayObject*)PyArray_SimpleNew(1, dims, NPY_OBJECT);
-        // Py_INCREF(array_from); // normalize refs when casting
     }
     else if (dtype_is_unicode || dtype_is_string) {
         array_to = (PyArrayObject*)PyArray_Zeros(1, dims, dtype, 0); // steals dtype ref
-        // Py_INCREF(array_from); // normalize refs when casting
     }
     else {
         array_to = (PyArrayObject*)PyArray_Empty(1, dims, dtype, 0); // steals dtype ref
-        // if (PyArray_TYPE(array_from) == NPY_DATETIME &&
-        //         PyArray_TYPE(array_to) == NPY_DATETIME &&
-        //         AK_dt_unit_from_array(array_from) != AK_dt_unit_from_array(array_to)
-        //         ) {
-        //     // if trying to cast into a dt64 array, need to pre-convert; array_from is originally borrowed; calling cast sets it to a new ref
-        //     dtype = PyArray_DESCR(array_to); // borrowed ref
-        //     Py_INCREF(dtype);
-        //     array_from = (PyArrayObject*)PyArray_CastToType(array_from, dtype, 0);
-        // }
-        // else {
-        //     Py_INCREF(array_from); // normalize refs when casting
-        // }
     }
     if (array_to == NULL) {
         PyErr_SetNone(PyExc_MemoryError);
-        // Py_DECREF((PyObject*)array_from);
         return NULL;
     }
 

test/test_tri_map.py

@@ -1274,3 +1274,42 @@ def test_tri_map_merge_g(self) -> None:
 
         post = tm.map_merge(src, dst)
         self.assertEqual(post.tolist(), [None, False, -42, 'ee', 'ff'])
+
+    def test_tri_map_merge_h(self) -> None:
+        src = np.array([0, 200, 300, 40], dtype=np.int64)
+        dst = np.array([0, 200, 300, 40, 50], dtype=np.int64)
+
+        tm = TriMap(len(src), len(dst))
+        tm.register_one(0, 0)
+        tm.register_one(1, 1)
+        tm.register_one(2, 2)
+        tm.register_one(3, 3)
+        tm.register_unmatched_dst()
+        tm.finalize()
+
+        post = tm.map_merge(src, dst)
+        self.assertEqual(post.tolist(), [0, 200, 300, 40, 50])
+
+    def test_tri_map_merge_i(self) -> None:
+        src = np.array(['2022-01', '1954-03', '1743-09', '1988-12'], dtype=np.datetime64)
+        dst = np.array(['1743-09', '2022-01', '2022-01', '2022-01', '1743-09', '2005-11'], dtype=np.datetime64)
+
+        tm = TriMap(len(src), len(dst))
+        tm.register_many(0, np.array([1, 2, 3], dtype=np.dtype(np.int64)))
+        tm.register_one(1, -1)
+        tm.register_many(2, np.array([0, 4], dtype=np.dtype(np.int64)))
+        tm.register_one(3, -1)
+        tm.register_unmatched_dst()
+        tm.finalize()
+
+        post = tm.map_merge(src, dst)
+        self.assertEqual(list(post),
+            [np.datetime64('2022-01'),
+             np.datetime64('2022-01'),
+             np.datetime64('2022-01'),
+             np.datetime64('1954-03'),
+             np.datetime64('1743-09'),
+             np.datetime64('1743-09'),
+             np.datetime64('1988-12'),
+             np.datetime64('2005-11')]
+             )