diff --git a/tools/releasetools/non_ab_ota.py b/tools/releasetools/non_ab_ota.py
index 80c3083..5839d17 100644
--- a/tools/releasetools/non_ab_ota.py
+++ b/tools/releasetools/non_ab_ota.py
@@ -26,7 +26,7 @@
 
 import common
 import edify_generator
-from edify_generator import ErrorCode, PARTITION_TYPES
+from edify_generator import ErrorCode
 from check_target_files_vintf import CheckVintfIfTrebleEnabled, HasPartition
 from common import OPTIONS, Run, MakeTempDir, RunAndCheckOutput, ZipWrite, MakeTempFile
 from ota_utils import UNZIP_PATTERN, FinalizeMetadata, GetPackageMetadata, PropertyFiles
@@ -60,9 +60,9 @@
     # Disable imgdiff because it relies on zlib to produce stable output
     # across different versions, which is often not the case.
     return BlockDifference(name, partition_tgt, partition_src,
-                           check_first_block,
-                           version=blockimgdiff_version,
-                           disable_imgdiff=True)
+                                  check_first_block,
+                                  version=blockimgdiff_version,
+                                  disable_imgdiff=True)
 
   if source_zip:
     # See notes in common.GetUserImage()
@@ -85,7 +85,7 @@
                                 info_dict=target_info,
                                 reset_file_map=True)
       block_diff_dict[partition] = BlockDifference(partition, tgt,
-                                                   src=None)
+                                                          src=None)
     # Incremental OTA update.
     else:
       block_diff_dict[partition] = GetIncrementalBlockDifferenceForPartition(
@@ -413,7 +413,7 @@
                           block_diff_dict.values()]
   if updating_boot:
     boot_type, boot_device_expr = GetTypeAndDeviceExpr("/boot",
-                                                       source_info)
+                                                              source_info)
     d = Difference(target_boot, source_boot, "bsdiff")
     _, _, d = d.ComputePatch()
     if d is None:
@@ -1544,6 +1544,16 @@
         append('move %s %s' % (p, u.tgt_group))
 
 
+# map recovery.fstab's fs_types to mount/format "partition types"
+PARTITION_TYPES = {
+    "ext4": "EMMC",
+    "emmc": "EMMC",
+    "f2fs": "EMMC",
+    "squashfs": "EMMC",
+    "erofs": "EMMC"
+}
+
+
 def GetTypeAndDevice(mount_point, info, check_no_slot=True):
   """
   Use GetTypeAndDeviceExpr whenever possible. This function is kept for