partition_table: Support same fallback logic as bootloader for defaul…

…t boot partition

Generates correct "make flash" command even when partition table has no factory partition.

Also adds unit tests for parttool.py

Closes #2086

components/partition_table/Makefile.projbuild

@@ -57,7 +57,7 @@ all_binaries: $(PARTITION_TABLE_BIN) partition_table_get_info
 
 partition_table_get_info: $(PARTITION_TABLE_BIN)
 	$(eval PHY_DATA_OFFSET:=$(shell $(GET_PART_INFO) --type data --subtype phy --offset $(PARTITION_TABLE_BIN)))
-	$(eval APP_OFFSET:=$(shell $(GET_PART_INFO) --type app --subtype factory --offset $(PARTITION_TABLE_BIN)))
+	$(eval APP_OFFSET:=$(shell $(GET_PART_INFO) --default-boot-partition --offset $(PARTITION_TABLE_BIN)))
 
 export APP_OFFSET
 export PHY_DATA_OFFSET

components/partition_table/gen_esp32part.py

@@ -106,6 +106,38 @@ def __getitem__(self, item):
         else:
             return super(PartitionTable, self).__getitem__(item)
 
+    def find_by_type(self, ptype, subtype):
+        """ Return a partition by type & subtype, returns
+        None if not found """
+        TYPES = PartitionDefinition.TYPES
+        SUBTYPES = PartitionDefinition.SUBTYPES
+        # convert ptype & subtypes names (if supplied this way) to integer values
+        try:
+            ptype = TYPES[ptype]
+        except KeyError:
+            try:
+                ptypes = int(ptype, 0)
+            except TypeError:
+                pass
+        try:
+            subtype = SUBTYPES[int(ptype)][subtype]
+        except KeyError:
+            try:
+                ptypes = int(ptype, 0)
+            except TypeError:
+                pass
+
+        for p in self:
+            if p.type == ptype and p.subtype == subtype:
+                return p
+        return None
+
+    def find_by_name(self, name):
+        for p in self:
+            if p.name == name:
+                return p
+        return None
+
     def verify(self):
         # verify each partition individually
         for p in self:
@@ -202,7 +234,7 @@ class PartitionDefinition(object):
         "encrypted" : 0
     }
 
-    # add subtypes for the 16 OTA slot values ("ota_XXX, etc.")
+    # add subtypes for the 16 OTA slot values ("ota_XX, etc.")
     for ota_slot in range(16):
         SUBTYPES[TYPES["app"]]["ota_%d" % ota_slot] = 0x10 + ota_slot
 

components/partition_table/parttool.py

@@ -48,17 +48,30 @@ def main():
     parser = argparse.ArgumentParser(description='Returns info about the required partition.')
 
     parser.add_argument('--quiet', '-q', help="Don't print status messages to stderr", action='store_true')
-    parser.add_argument('--partition-name', '-p', help='The name of the required partition', type=str, default=None)
-    parser.add_argument('--type', '-t', help='The type of the required partition', type=str, default=None)
+
+    search_type = parser.add_mutually_exclusive_group()
+    search_type.add_argument('--partition-name', '-p', help='The name of the required partition', type=str, default=None)
+    search_type.add_argument('--type', '-t', help='The type of the required partition', type=str, default=None)
+    search_type.add_argument('--default-boot-partition', help='Select the default boot partition, '+
+                             'using the same fallback logic as the IDF bootloader', action="store_true")
+
     parser.add_argument('--subtype', '-s', help='The subtype of the required partition', type=str, default=None)
-
-    parser.add_argument('--offset', '-o', help='Return offset of required partition', action="store_true", default=None)
-    parser.add_argument('--size', help='Return size of required partition', action="store_true", default=None)
-
-    parser.add_argument('input', help='Path to CSV or binary file to parse. Will use stdin if omitted.', type=argparse.FileType('rb'), default=sys.stdin)
+
+    parser.add_argument('--offset', '-o', help='Return offset of required partition', action="store_true")
+    parser.add_argument('--size', help='Return size of required partition', action="store_true")
+
+    parser.add_argument('input', help='Path to CSV or binary file to parse. Will use stdin if omitted.',
+                        type=argparse.FileType('rb'), default=sys.stdin)
 
     args = parser.parse_args()
 
+    if args.type is not None and args.subtype is None:
+        status("If --type is specified, --subtype is required")
+        return 2
+    if args.type is None and args.subtype is not None:
+        status("--subtype is only used with --type")
+        return 2
+
     quiet = args.quiet
 
     input = args.input.read()
@@ -71,36 +84,30 @@ def main():
         status("Parsing CSV input...")
         table = gen.PartitionTable.from_csv(input)
 
-    if args.partition_name is not None:
-        offset = 0
-        size = 0
-        for p in table:
-            if p.name == args.partition_name:
-                offset = p.offset
-                size = p.size
-                break;
-        if args.offset is not None:
-            print('0x%x ' % (offset))
-        if args.size is not None:
-            print('0x%x' % (size))
-        return 0
-
-    if args.type is not None and args.subtype is not None:
-        offset = 0
-        size = 0
-        TYPES = gen.PartitionDefinition.TYPES
-        SUBTYPES = gen.PartitionDefinition.SUBTYPES
-        for p in table:
-            if p.type == TYPES[args.type]:
-                if p.subtype == SUBTYPES[TYPES[args.type]][args.subtype]:
-                    offset = p.offset
-                    size = p.size
-                    break;
-        if args.offset is not None:
-            print('0x%x ' % (offset))
-        if args.size is not None:
-            print('0x%x' % (size))
-        return 0
+    found_partition = None
+
+    if args.default_boot_partition:
+        search = [ "factory" ] + [ "ota_%d" % d for d in range(16) ]
+        for subtype in search:
+            found_partition = table.find_by_type("app", subtype)
+            if found_partition is not None:
+                break
+    elif args.partition_name is not None:
+        found_partition = table.find_by_name(args.partition_name)
+    elif args.type is not None:
+        found_partition = table.find_by_type(args.type, args.subtype)
+    else:
+        raise RuntimeError("invalid partition selection choice")
+
+    if found_partition is None:
+        return 1  # nothing found
+
+    if args.offset:
+        print('0x%x ' % (found_partition.offset))
+    if args.size:
+        print('0x%x' % (found_partition.size))
+
+    return 0
 
 class InputError(RuntimeError):
     def __init__(self, e):
@@ -115,7 +122,8 @@ def __init__(self, partition, message):
 
 if __name__ == '__main__':
     try:
-        main()
+        r = main()
+        sys.exit(r)
     except InputError as e:
         print(e, file=sys.stderr)
         sys.exit(2)

components/partition_table/test_gen_esp32part_host/gen_esp32part_tests.py

@@ -356,5 +356,53 @@ def test_bad_alignment(self):
             t.verify()
 
 
+class PartToolTests(unittest.TestCase):
+
+    def _run_parttool(self, csvcontents, args):
+        csvpath = tempfile.mktemp()
+        with open(csvpath, "w") as f:
+            f.write(csvcontents)
+        try:
+            return subprocess.check_output([sys.executable, "../parttool.py"] + args.split(" ") + [ csvpath ]).strip()
+        finally:
+            os.remove(csvpath)
+
+    def test_find_basic(self):
+        csv = """
+nvs,      data, nvs,     0x9000,  0x4000
+otadata,  data, ota,     0xd000,  0x2000
+phy_init, data, phy,     0xf000,  0x1000
+factory,  app, factory, 0x10000,  1M
+        """
+        rpt = lambda args: self._run_parttool(csv, args)
+
+        self.assertEqual(
+            rpt("--type data --subtype nvs --offset"), "0x9000")
+        self.assertEqual(
+            rpt("--type data --subtype nvs --size"), "0x4000")
+        self.assertEqual(
+            rpt("--partition-name otadata --offset"), "0xd000")
+        self.assertEqual(
+            rpt("--default-boot-partition --offset"), "0x10000")
+
+    def test_fallback(self):
+        csv = """
+nvs,      data, nvs,     0x9000,  0x4000
+otadata,  data, ota,     0xd000,  0x2000
+phy_init, data, phy,     0xf000,  0x1000
+ota_0,  app,    ota_0,   0x30000,  1M
+ota_1,  app,    ota_1,          ,  1M
+        """
+        rpt = lambda args: self._run_parttool(csv, args)
+
+        self.assertEqual(
+            rpt("--type app --subtype ota_1 --offset"), "0x130000")
+        self.assertEqual(
+            rpt("--default-boot-partition --offset"), "0x30000")  # ota_0
+        csv_mod = csv.replace("ota_0", "ota_2")
+        self.assertEqual(
+            self._run_parttool(csv_mod, "--default-boot-partition --offset"),
+            "0x130000")  # now default is ota_1
+
 if __name__ =="__main__":
     unittest.main()