diff --git a/androscalpel/src/apk.rs b/androscalpel/src/apk.rs
index 8c48bc8..3cfe6ae 100644
--- a/androscalpel/src/apk.rs
+++ b/androscalpel/src/apk.rs
@@ -3284,18 +3284,27 @@ impl Apk {
     }
 
     #[cfg_attr(feature = "python", pyo3(signature = (class, dex_file=None)))]
-    pub fn remove_class(&mut self, class: &IdType, dex_file: Option<&str>) -> Result<()> {
+    pub fn remove_class(
+        &mut self,
+        class: &IdType,
+        dex_file: Option<&str>,
+    ) -> Result<Option<Class>> {
         if let Some(dex_file) = dex_file {
-            self.dex_files
+            Ok(self
+                .dex_files
                 .get_mut(dex_file)
                 .with_context(|| format!("file {} not found in apk", dex_file))?
                 .classes
-                .remove(class);
+                .remove(class))
         } else {
-            for DexFile { classes, .. } in self.dex_files.values_mut() {
-                classes.remove(class);
-            }
+            // TODO valid_file_only ?
+            Ok(self.get_file_of_class(class, false).and_then(|file| {
+                self.dex_files
+                    .get_mut(&file)
+                    .expect("`get_file_of_class()` return a string not in `dex_files`")
+                    .classes
+                    .remove(class)
+            }))
         }
-        Ok(())
     }
 }