size_report: use anytree

Use anytree module to display tree and to allow easy exporting into
json.
Add option to export results into a json file.

Signed-off-by: Anas Nashif <anas.nashif@intel.com>

scripts/footprint/size_report

@@ -9,7 +9,6 @@
 
 import argparse
 import os
-import platform
 import sys
 import re
 from pathlib import Path
@@ -28,6 +27,11 @@ from elftools.dwarf.locationlists import (
 if LooseVersion(elftools.__version__) < LooseVersion('0.24'):
     sys.exit("pyelftools is out of date, need version 0.24 or later")
 
+from colorama import Fore
+
+from anytree import RenderTree, NodeMixin, findall_by_attr
+from anytree.exporter import JsonExporter
+
 
 # ELF section flags
 SHF_WRITE = 0x1
@@ -106,29 +110,6 @@ def match_symbol_address(symlist, die, parser, dwarfinfo):
     return None
 
 
-def parse_args():
-    global args
-
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument("-k", "--kernel", required=True,
-                        help="Zephyr ELF binary")
-    parser.add_argument("-z", "--zephyrbase", required=True,
-                        help="Zephyr base path")
-    parser.add_argument("-o", "--output", required=True,
-                        help="Output path")
-    parser.add_argument("target", choices=['rom', 'ram'])
-    parser.add_argument("-d", "--depth", dest="depth", type=int,
-                        help="How deep should we go into the tree",
-                        metavar="DEPTH")
-    parser.add_argument("-v", "--verbose", action="store_true",
-                        help="Print extra debugging information")
-    parser.add_argument("--nocolor", action="store_true",
-                        help="No color output")
-
-    args = parser.parse_args()
-
-
 # Fetch the symbols from the symbol table and put them
 # into ROM, RAM buckets
 def get_symbols(elf, addr_ranges):
@@ -495,89 +476,104 @@ def set_root_path_for_unmapped_symbols(symbol_dict, addr_range, processed):
     processed['mapped_addr'] = mapped_addresses
     processed['unmapped_symbols'] = unmapped_symbols
 
+class TreeNode(NodeMixin):
+    def __init__(self, name, identifier, size=0, parent=None, children=None):
+        super(TreeNode, self).__init__()
+        self.name = name
+        self.size = size
+        self.parent = parent
+        self.identifier = identifier
+        if children:
+            self.children = children
 
-def generate_tree(symbol_dict):
+    def __repr__(self):
+        return self.name
+
+def generate_any_tree(symbol_dict):
+    root = TreeNode('root', ":")
 
     # A set of helper function for building a simple tree with a path-like
     # hierarchy.
-    def _insert_one_elem(tree, path, size):
+    def _insert_one_elem(root, path, size):
         cur = None
+        node = None
+        parent = root
         for part in path.parts:
             if cur is None:
                 cur = part
             else:
                 cur = str(Path(cur, part))
-            if cur in tree:
-                tree[cur] += size
-            else:
-                tree[cur] = size
 
-    total_size = 0
+            results = findall_by_attr(root, cur, name="identifier")
-    nodes = {}
+            if results:
-    nodes[':'] = 0
+                item = results[0]
+                item.size += size
+                parent = item
+            else:
+                if node:
+                    parent = node
+                node = TreeNode(name=str(part), identifier=cur, size=size, parent=parent)
 
     for name, sym in symbol_dict.items():
         for symbol in sym:
             size = get_symbol_size(symbol['symbol'])
             for file in symbol['mapped_files']:
                 path = Path(file, name)
-                _insert_one_elem(nodes, path, size)
+                if path.is_absolute():
+                    zb = Path(args.zephyrbase)
+                    if zb.parent in path.parents:
+                        path = path.relative_to(zb.parent)
 
-    ret = {'nodes': nodes,
+                _insert_one_elem(root, path, size)
-           'size': total_size}
-    return ret
 
+    return root
 
-def print_tree(data, total, depth):
+def node_sort(items):
-    base = args.zephyrbase
+    return sorted(items, key=lambda item: item.name)
-    totp = 0
 
-    bcolors_ansi = {
+def print_any_tree(root, total_size):
-        "HEADER"    : '\033[95m',
-        "OKBLUE"    : '\033[94m',
-        "OKGREEN"   : '\033[92m',
-        "WARNING"   : '\033[93m',
-        "FAIL"      : '\033[91m',
-        "ENDC"      : '\033[0m',
-        "BOLD"      : '\033[1m',
-        "UNDERLINE" : '\033[4m'
-    }
-    if platform.system() == "Windows" or args.nocolor:
-        # Set all color codes to empty string on Windows
-        #
-        # TODO: Use an approach like the pip package 'colorama' to
-        # support colors on Windows
-        bcolors = dict.fromkeys(bcolors_ansi, '')
-    else:
-        bcolors = bcolors_ansi
 
-    print('{:92s} {:10s} {:8s}'.format(
+    print('{:101s} {:7s} {:8s}'.format(
-        bcolors["FAIL"] + "Path", "Size", "%" + bcolors["ENDC"]))
+        Fore.YELLOW + "Path", "Size", "%" + Fore.RESET))
     print('=' * 110)
-    for i in sorted(data):
+    for row in RenderTree(root, childiter=node_sort):
-        p = i.split(os.path.sep)
+        f = len(row.pre) + len(row.node.name)
-        if depth and len(p) > depth:
+        s = str(row.node.size).rjust(100-f)
-            continue
+        percent = 100 * float(row.node.size) / float(total_size)
 
-        percent = 100 * float(data[i]) / float(total)
+        cc = cr = ""
-        percent_c = percent
+        if not row.node.children:
-        if len(p) < 2:
+            cc = Fore.CYAN
-            totp += percent
+            cr = Fore.RESET
-
+        elif row.node.name.endswith(".c") or row.node.name.endswith(".h"):
-        if len(p) > 1:
+            cc = Fore.GREEN
-            if not os.path.exists(os.path.join(base, i)):
+            cr = Fore.RESET
-                s = bcolors["WARNING"] + p[-1] + bcolors["ENDC"]
-            else:
-                s = bcolors["OKBLUE"] + p[-1] + bcolors["ENDC"]
-            print('{:80s} {:20d} {:8.2f}%'.format(
-                "  " * (len(p) - 1) + s, data[i], percent_c))
-        else:
-            print('{:80s} {:20d} {:8.2f}%'.format(
-                bcolors["OKBLUE"] + i + bcolors["ENDC"], data[i], percent_c))
 
+        print(f"{row.pre}{cc}{row.node.name}{cr} {s} {Fore.BLUE}{percent:.2f}%{Fore.RESET}")
     print('=' * 110)
-    print('{:92d}'.format(total))
+    print(f'{total_size:>101}')
-    return totp
+
+
+
+def parse_args():
+    global args
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("-k", "--kernel", required=True,
+                        help="Zephyr ELF binary")
+    parser.add_argument("-z", "--zephyrbase", required=True,
+                        help="Zephyr base path")
+    parser.add_argument("-o", "--output", required=True,
+                        help="Output path")
+    parser.add_argument("target", choices=['rom', 'ram'])
+    parser.add_argument("-d", "--depth", dest="depth", type=int,
+                        help="How deep should we go into the tree",
+                        metavar="DEPTH")
+    parser.add_argument("-v", "--verbose", action="store_true",
+                        help="Print extra debugging information")
+    parser.add_argument("--json", help="store results in a JSON file.")
+    args = parser.parse_args()
 
 
 def main():
@@ -623,8 +619,15 @@ def main():
             for sym in processed['unmapped_symbols']:
                 print("INFO: Unmapped symbol: {0}".format(sym))
 
-        tree = generate_tree(symbol_dict)
+        root = generate_any_tree(symbol_dict)
-        print_tree(tree['nodes'], symsize, args.depth)
+        print_any_tree(root, symsize)
+
+        if args.json:
+            exporter = JsonExporter(indent=2, sort_keys=True)
+            data = exporter.export(root)
+            with open(args.json, "w") as fp:
+                fp.write(data)
+
 
 if __name__ == "__main__":
     main()