From d1ca44450d8cbd1f8a9ba15586eb1582677d5648 Mon Sep 17 00:00:00 2001
From: Maksim Andrianov <maksimandrianov1@gmail.com>
Date: Tue, 26 May 2020 15:34:40 +0300
Subject: [PATCH] [python][generator] Added check categories.

---
 tools/python/maps_generator/checks/check.py   | 15 ++++--
 .../maps_generator/checks/check_categories.py | 54 +++++++++++++++++++
 2 files changed, 64 insertions(+), 5 deletions(-)
 create mode 100644 tools/python/maps_generator/checks/check_categories.py

diff --git a/tools/python/maps_generator/checks/check.py b/tools/python/maps_generator/checks/check.py
index b686ad5dcf..5d43c3c5ca 100644
--- a/tools/python/maps_generator/checks/check.py
+++ b/tools/python/maps_generator/checks/check.py
@@ -3,6 +3,7 @@ from abc import ABC
 from abc import abstractmethod
 from collections import namedtuple
 from enum import Enum
+from functools import lru_cache
 from typing import Any
 from typing import Callable
 from typing import List
@@ -208,6 +209,14 @@ class CompareCheckSet(Check):
         return (c for c in self.checks if c.get_result() is not None)
 
 
+@lru_cache(maxsize=None)
+def _get_and_check_files(old_path, new_path, ext):
+    files = list(filter(lambda f: f.endswith(ext), os.listdir(old_path)))
+    s = set(files) ^ set(filter(lambda f: f.endswith(ext), os.listdir(new_path)))
+    assert len(s) == 0, s
+    return files
+
+
 def build_check_set_for_files(
     name: str,
     old_path: str,
@@ -226,12 +235,8 @@ def build_check_set_for_files(
             f"CheckSetBuilderForFiles is not implemented for recursive."
         )
 
-    files = list(filter(lambda f: f.endswith(ext), os.listdir(old_path)))
-    s = set(files) ^ set(filter(lambda f: f.endswith(ext), os.listdir(new_path)))
-    assert len(s) == 0, s
-
     cs = CompareCheckSet(name)
-    for file in files:
+    for file in _get_and_check_files(old_path, new_path, ext):
         cs.add_check(
             CompareCheck(
                 file, os.path.join(old_path, file), os.path.join(new_path, file)
diff --git a/tools/python/maps_generator/checks/check_categories.py b/tools/python/maps_generator/checks/check_categories.py
new file mode 100644
index 0000000000..a8b87754ff
--- /dev/null
+++ b/tools/python/maps_generator/checks/check_categories.py
@@ -0,0 +1,54 @@
+from collections import defaultdict
+
+from maps_generator.checks import check
+from mwm.types import NAME_TO_INDEX_TYPE_MAPPING
+from maps_generator.checks.check_mwm_types import count_all_types
+
+
+def parse_groups(path):
+    groups = defaultdict(set)
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if line.startswith("#"):
+                continue
+
+            if line.startswith("@"):
+                continue
+
+            array = line.split("@", maxsplit=1)
+            if len(array) == 2:
+                types_str, categories = array
+                types_int = {
+                    NAME_TO_INDEX_TYPE_MAPPING[t]
+                    for t in types_str.strip("|").split("|")
+                }
+                for category in categories.split("|"):
+                    category = category.replace("@", "", 1)
+                    groups[category].update(types_int)
+        return groups
+
+
+def get_mwm_categories_check_set(
+    old_path: str, new_path: str, categories_path: str
+) -> check.CompareCheckSet:
+    cs = check.CompareCheckSet("Sections categories check")
+
+    def make_do(indexes):
+        def do(path):
+            all_types = count_all_types(path)
+            return sum(all_types[i] for i in indexes)
+
+        return do
+
+    for category, types in parse_groups(categories_path).items():
+        cs.add_check(
+            check.build_check_set_for_files(
+                f"Category {category} check",
+                old_path,
+                new_path,
+                ext=".mwm",
+                do=make_do(types),
+            )
+        )
+    return cs