refactor validate code and add tests for set_code and set_level

Author: Jiang Yue

doc/source/whatsnew/v0.25.0.rst

@@ -387,8 +387,7 @@ MultiIndex
 ^^^^^^^^^^
 
 - Bug in which incorrect exception raised by :class:`Timedelta` when testing the membership of :class:`MultiIndex` (:issue:`24570`)
-- Bug in :class:`Multindex` construction from levels and codes that would incorrectly allows code values < -1
-- Bug in :class:`Multindex` construction from levels and codes that would incorrectly allows NaN levels (:issue:`19387`)
+- Bug in :class:`Multindex` construction from levels and codes that would incorrectly allows code values < -1 or NaN levels (:issue:`19387`)
 
 I/O
 ^^^

pandas/core/indexes/multi.py

@@ -244,17 +244,12 @@ def __new__(cls, levels=None, codes=None, sortorder=None, names=None,
         if verify_integrity:
             result._verify_integrity()
 
-        codes = [cls._reassign_na_codes(level, code)
-                 for level, code in zip(levels, codes)]
-        result._set_codes(codes, validate=False)
-
         if _set_identity:
             result._reset_identity()
 
         return result
 
-    @classmethod
-    def _reassign_na_codes(cls, level, code):
+    def _validate_codes(cls, level, code):
         null_mask = isna(level)
         if np.any(null_mask):
             code = np.where(null_mask[code], -1, code)
@@ -290,19 +285,26 @@ def _verify_integrity(self, codes=None, levels=None):
                 raise ValueError("Unequal code lengths: %s" %
                                  ([len(code_) for code_ in codes]))
             if len(level_codes) and level_codes.max() >= len(level):
-                raise ValueError("On level %d, code max (%d) >= length of"
-                                 " level  (%d). NOTE: this index is in an"
-                                 " inconsistent state" % (i, level_codes.max(),
-                                                          len(level)))
+                raise ValueError("On level {level}, code max ({max_code})"
+                                 " >= length of level  ({level_len}). "
+                                 "NOTE: this index is in an inconsistent"
+                                 " state".format(
+                                     level=i, max_code=level_codes.max(),
+                                     level_len=len(level)))
             if len(level_codes) and level_codes.min() < -1:
-                raise ValueError("On level %d, code value (%d) < -1" %
-                                 (i, level_codes.min()))
+                raise ValueError("On level {level}, code value ({code})"
+                                 " < -1".format(
+                                     level=i, code=level_codes.min()))
             if not level.is_unique:
                 raise ValueError("Level values must be unique: {values} on "
                                  "level {level}".format(
                                      values=[value for value in level],
                                      level=i))
 
+        codes = [self._validate_codes(level, code)
+                 for level, code in zip(levels, codes)]
+        self._set_codes(codes, validate=False)
+
     @classmethod
     def from_arrays(cls, arrays, sortorder=None, names=None):
         """
@@ -691,7 +693,6 @@ def labels(self):
 
     def _set_codes(self, codes, level=None, copy=False, validate=True,
                    verify_integrity=False):
-
         if validate and level is None and len(codes) != self.nlevels:
             raise ValueError("Length of codes must match number of levels")
         if validate and level is not None and len(codes) != len(level):
@@ -712,8 +713,9 @@ def _set_codes(self, codes, level=None, copy=False, validate=True,
 
         if verify_integrity:
             self._verify_integrity(codes=new_codes)
+        else:
+            self._codes = new_codes
 
-        self._codes = new_codes
         self._tuples = None
         self._reset_cache()
 

pandas/tests/indexes/multi/test_constructor.py

@@ -90,16 +90,30 @@ def test_constructor_mismatched_codes_levels(idx):
 
 
 def test_na_levels():
-    tm.assert_index_equal(
-        MultiIndex(levels=[[np.nan, None, pd.NaT, 128, 2]],
-                   codes=[[0, -1, 1, 2, 3, 4]]),
-        MultiIndex(levels=[[np.nan, None, pd.NaT, 128, 2]],
-                   codes=[[-1, -1, -1, -1, 3, 4]]))
-    tm.assert_index_equal(
-        MultiIndex(levels=[[np.nan, 's', pd.NaT, 128, None]],
-                   codes=[[0, -1, 1, 2, 3, 4]]),
-        MultiIndex(levels=[[np.nan, 's', pd.NaT, 128, None]],
-                   codes=[[-1, -1, 1, -1, 3, -1]]))
+    # GH26408
+    result = MultiIndex(levels=[[np.nan, None, pd.NaT, 128, 2]],
+                        codes=[[0, -1, 1, 2, 3, 4]])
+    expected = MultiIndex(levels=[[np.nan, None, pd.NaT, 128, 2]],
+                          codes=[[-1, -1, -1, -1, 3, 4]])
+    tm.assert_index_equal(result, expected)
+
+    result = MultiIndex(levels=[[np.nan, 's', pd.NaT, 128, None]],
+                        codes=[[0, -1, 1, 2, 3, 4]])
+    expected = MultiIndex(levels=[[np.nan, 's', pd.NaT, 128, None]],
+                          codes=[[-1, -1, 1, -1, 3, -1]])
+    tm.assert_index_equal(result, expected)
+
+    # verify set_levels and set_codes
+    result = MultiIndex(
+        levels=[[1, 2, 3, 4, 5]], codes=[[0, -1, 1, 2, 3, 4]]).set_levels(
+            [[np.nan, 's', pd.NaT, 128, None]])
+    tm.assert_index_equal(result, expected)
+
+    result = MultiIndex(
+        levels=[[np.nan, 's', pd.NaT, 128, None]],
+        codes=[[1, 2, 2, 2, 2, 2]]).set_codes(
+            [[0, -1, 1, 2, 3, 4]])
+    tm.assert_index_equal(result, expected)
 
 
 def test_labels_deprecated(idx):