Update map_processor

scripts/README.md

@@ -6,19 +6,43 @@ This script can:
 
 To use this script it have to be in the same folder as [integration files](../custom_components/xiaomi_cloud_map_extractor).
 
-Dependencies installation:
+Config file (passed with argument `--config`) should contain configuration of `xiaomi_cloud_map_extractor` camera (cloud credentials are used only for `download` mode).
+
+### Dependencies installation
 ```bash
 pip3 install pillow pybase64 python requests pycryptodome pyyaml
 ```
 
-Downloading map from Xiaomi cloud:
+### Downloading map from Xiaomi cloud
 ```bash
 python3 map_processor.py download --config camera.yaml
 ```
 
-Parsing already downloaded raw map file:
+### Parsing already downloaded raw map file
 ```bash
 python3 map_processor.py parse --config camera.yaml --map-file map_data.gz --api xiaomi
 ```
 
-Config file (passed with argument `--config`) should contain configuration of `xiaomi_cloud_map_extractor` camera.
+### Testing multiple raw map files
+```bash
+python3 map_processor.py test --config camera.yaml --test-data test_data
+```
+Supported file structure:
+```
+|-- camera.yaml
+`-- test_data
+    |-- dreame
+    |   |-- map_data_dreame.vacuum.mc1808.b64
+    |   |-- map_data_dreame.vacuum.p2028.b64
+    |   `-- map_data_dreame.vacuum.p2259.b64
+    |-- roidmi
+    |   |-- map_data_roidmi.vacuum.v60_1.gz
+    |   `-- map_data_viomi.vacuum.v18_1.gz
+    |-- viomi
+    |   |-- map_data_viomi.vacuum.v6.zlib
+    |   `-- map_data_viomi.vacuum.v7_1.zlib
+    `-- xiaomi
+        |-- map_data_roborock.vacuum.a08.gz
+        |-- map_data_roborock.vacuum.s5.gz
+        `-- map_data_rockrobo.vacuum.v1.gz
+```

scripts/map_processor.py

@@ -64,7 +64,7 @@ def attributes_to_dict(attributes):
     return attributes
 
 
-def parse_map_file(map_config, map_filename, api):
+def parse_map_file(map_config, map_filename, api, suffix=""):
     print(f"Parsing map file \"{map_filename}\" with api \"{api}\"")
     map_file = open(map_filename, "rb").read()
     colors = map_config[CONF_COLORS]
@@ -79,19 +79,22 @@ def parse_map_file(map_config, map_filename, api):
         drawables = CONF_AVAILABLE_DRAWABLES[1:]
 
     map_data = None
-    if api == "xiaomi":
-        map_data = XiaomiVacuum.decode_map(None, map_file, colors, drawables, texts, sizes, transform)
-    elif api == "viomi":
-        map_data = ViomiVacuum.decode_map(None, map_file, colors, drawables, texts, sizes, transform)
-    elif api == "roidmi":
-        map_data = RoidmiVacuum.decode_map(None, map_file, colors, drawables, texts, sizes, transform)
+    try:
+        if api == CONF_AVAILABLE_API_XIAOMI:
+            map_data = XiaomiVacuum.decode_map(None, map_file, colors, drawables, texts, sizes, transform)
+        elif api == CONF_AVAILABLE_API_VIOMI:
+            map_data = ViomiVacuum.decode_map(None, map_file, colors, drawables, texts, sizes, transform)
+        elif api == CONF_AVAILABLE_API_ROIDMI:
+            map_data = RoidmiVacuum.decode_map(None, map_file, colors, drawables, texts, sizes, transform)
+    except Exception as e:
+        print(f"Failed to parse map data! {e}")
     if map_data is not None:
-        map_data.image.data.save(f"{map_filename}.png")
-        print(f"Map image saved to \"{map_filename}.png\"")
-        attributes_output_file = open(f"{map_filename}.yaml", "w")
+        map_data.image.data.save(f"{map_filename}{suffix}.png")
+        print(f"Map image saved to \"{map_filename}{suffix}.png\"")
+        attributes_output_file = open(f"{map_filename}{suffix}.yaml", "w")
         yaml.dump(attributes_to_dict(map_data), attributes_output_file)
         attributes_output_file.close()
-        print(f"Map attributes saved to \"{map_filename}.yaml\"")
+        print(f"Map attributes saved to \"{map_filename}{suffix}.yaml\"")
     else:
         print("Failed to parse map data!")
 
@@ -108,11 +111,27 @@ def run_download(map_config, data_output_dir):
     print(f"Map data successfully saved to \"{data_output_dir}\" directory!")
 
 
+def run_test(map_config, test_dir):
+    print("Running tests")
+    for api in list(filter(lambda d: d in CONF_AVAILABLE_APIS, os.listdir(test_dir))):
+        print(api)
+        for file in filter(lambda ff: os.path.isfile(ff),
+                           map(lambda f: f"{test_dir}/{api}/{f}", os.listdir(f"{test_dir}/{api}"))):
+            print("  " + file)
+            output = file + "_output"
+            if not os.path.exists(output):
+                os.mkdir(output)
+            parse_map_file(map_config, file, api, "_output/data")
+
+
 if __name__ == '__main__':
     args_parser = argparse.ArgumentParser(description='Map processor')
     args_subparsers = args_parser.add_subparsers(help="Available run modes", dest="mode")
     args_parser_download = args_subparsers.add_parser("download", help="Download and parse map")
     args_parser_download.add_argument("--config", type=str, required=True, help="camera yaml config file")
+    args_parser_test = args_subparsers.add_parser("test", help="Test multiple raw map files")
+    args_parser_test.add_argument("--config", type=str, required=True, help="camera yaml config file")
+    args_parser_test.add_argument("--test-data", type=str, required=True, help="test data directory")
     args_parser_parse = args_subparsers.add_parser("parse", help="Parse already downloaded map file")
     args_parser_parse.add_argument("--config", type=str, required=True, help="camera yaml config file")
     args_parser_parse.add_argument("--map-file", type=str, required=True, help="raw map file")
@@ -131,3 +150,5 @@ def run_download(map_config, data_output_dir):
         run_download(config, output_dir)
     elif args.mode == "parse":
         parse_map_file(config, args.map_file, args.api)
+    elif args.mode == "test":
+        run_test(config, args.test_data)