[ENH] MultiIndex: astype dict

pandas/core/frame.py

@@ -5529,8 +5529,7 @@ def set_index(
             else:
                 arrays.append(frame[col]._values)
                 names.append(col)
-                if drop:
-                    to_remove.append(col)
+                to_remove.append(col)
 
             if len(arrays[-1]) != len(self):
                 # check newest element against length of calling frame, since
@@ -5547,14 +5546,28 @@ def set_index(
             raise ValueError(f"Index has duplicate keys: {duplicates}")
 
         # use set to handle duplicate column names gracefully in case of drop
-        for c in set(to_remove):
-            del frame[c]
-
+        if len(to_remove) > 0:
+            to_remove = set(to_remove)
+            cast = frame.dtypes.filter(to_remove, axis=0).to_dict()
+            if drop:
+                for c in to_remove:
+                    del frame[c]
+        else:
+            cast = None
+
         # clear up memory usage
         index._cleanup()
-
+
+        # Ensure Index typing
+        if cast is not None:
+            if isinstance(self.index, MultiIndex):
+                index = index.astype(cast, copy=False)
+            else:
+                for _ in cast.values():
+                    index = index.astype(_, copy=False)
+
         frame.index = index
-
+        
         if not inplace:
             return frame
 

pandas/core/indexes/multi.py

@@ -3669,8 +3669,36 @@ def _convert_can_do_setop(self, other):
     # --------------------------------------------------------------------
 
     @doc(Index.astype)
-    def astype(self, dtype, copy: bool = True):
-        dtype = pandas_dtype(dtype)
+    def astype(self, dtype, copy: bool = True) -> MultiIndex:
+        try:
+            dtype = pandas_dtype(dtype)
+        except Exception as e:
+            # dict typing with index level names/int
+            if not isinstance(dtype, dict):
+                try:
+                    dtype = dict(dtype)
+                except Exception:
+                    dtype = dict(zip(range(len(self.levels)), dtype))
+
+            if isinstance(dtype, dict):
+                try:
+                    dtype = {
+                        self.names.index(k): pandas_dtype(dtype[k]) for k in dtype
+                    }
+                except Exception:
+                    dtype = {
+                        k: pandas_dtype(dtype[k]) for k in dtype
+                    }
+                return self.set_levels(
+                    [
+                        self.levels[i].astype(dtype[i], copy=copy)
+                        for i in range(len(self.levels)) if i in dtype
+                    ],
+                    level=dtype.keys(), inplace=copy
+                )
+            else:
+                raise e
+
         if is_categorical_dtype(dtype):
             msg = "> 1 ndim Categorical are not supported at this time"
             raise NotImplementedError(msg)