Add function token count comments (rough draft)

Author: sparr

pico_process.py

@@ -133,17 +133,18 @@ def __str__(m):
 def print_token_count(num_tokens, **kwargs):
     print_size("tokens", num_tokens, 8192, **kwargs)
 
-def process_code(ctxt, source, input_count=False, count=False, lint=False, minify=False, rename=False, fail=True, want_count=True):
+def process_code(ctxt, source, input_count=False, count=False, lint=False, minify=False, rename=False, fail=True, want_count=True, annotate=False):
     need_lint = lint not in (None, False)
     need_minify = minify not in (None, False)
     need_rename = rename not in (None, False)
+    need_annotate = annotate not in (None, False)
 
-    if not need_lint and not need_minify and not (want_count and (count or input_count)):
+    if not need_lint and not need_minify and not (want_count and (count or input_count)) and not need_annotate:
         return True, ()
 
     ok = False
     tokens, errors = tokenize(source, ctxt)
-    if not errors and (need_lint or need_minify):
+    if not errors and (need_lint or need_minify or need_annotate):
         root, errors = parse(source, tokens)
 
     if not errors:
@@ -164,10 +165,30 @@ def process_code(ctxt, source, input_count=False, count=False, lint=False, minif
         if count:
             print_token_count(count_tokens(tokens), handler=count)
 
+        if annotate:
+            annotate_code(ctxt, source, root)
+
     if fail and errors:
         raise Exception("\n".join(map(str, errors)))
     return ok, errors
 
+# Calls function for the descendants of node, then for node
+def apply_node_tree(node, func):
+    if hasattr(node, "extra_children"):
+        for child in reversed(node.extra_children):
+            apply_node_tree(child, func)
+    for child in reversed(node.children):
+        apply_node_tree(child, func)
+    func(node)
+
+def annotate_code(ctxt, source, root):
+    def comment_before_function(node, source=source):
+        if node.type==NodeType.function:
+            tokens, errors = tokenize(PicoSource("temp", source.text[node.idx:node.endidx]), ctxt)
+            count = count_tokens(tokens)
+            source.text = f'{source.text[:node.idx]}-- token count: {count}\n{source.text[node.idx:]}'
+    apply_node_tree(root, comment_before_function)
+
 def echo_code(code, echo=True):
     code = from_pico_chars(code)
     if echo == True:
@@ -181,6 +202,7 @@ def echo_code(code, echo=True):
 from pico_lint import lint_code
 from pico_minify import minify_code
 from pico_rename import rename_tokens
+from pico_parse import NodeType
 
 # re-export some things for examples/etc.
 from pico_tokenize import is_identifier, is_ident_char, CustomPreprocessor

shrinko8.py

@@ -49,6 +49,9 @@ def ParsableCountHandler(prefix, name, size, limit):
 pgroup.add_argument("--no-count-compress", action="store_true", help="do not compress the cart just to print the compressed size")
 pgroup.add_argument("--no-count-tokenize", action="store_true", help="do not tokenize the cart just to print the token count")
 
+pgroup = parser.add_argument_group("annotate options")
+pgroup.add_argument("-a", "--annotate", action="store_true", help="enable annotating source during output")
+
 pgroup = parser.add_argument_group("script options")
 pgroup.add_argument("-s", "--script", help="manipulate the cart via a custom python script - see README for api details")
 pgroup.add_argument("--script-args", nargs=argparse.REMAINDER, help="send arguments directly to --script", default=())
@@ -178,7 +181,7 @@ def fail(msg):
 
     ok, errors = process_code(ctxt, src, input_count=args.input_count, count=args.count,
                               lint=args.lint, minify=args.minify, rename=args.rename,
-                              fail=False, want_count=not args.no_count_tokenize)
+                              fail=False, want_count=not args.no_count_tokenize, annotate=args.annotate)
     if errors:
         print("Lint errors:" if ok else "Compilation errors:")
         for error in errors: