Add function token count comments

Author: sparr

pico_annotate.py

@@ -0,0 +1,33 @@
+from textwrap import indent
+
+# 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, fail=True):
+    def comment_before_function(node, source=source):
+        if node.type==NodeType.function:
+            tokens, errors = tokenize(PicoSource("temp", source.text[node.idx:node.endidx]), ctxt)
+            if fail and errors:
+                raise Exception("\n".join(map(str, errors)))
+            token_count = count_tokens(tokens)
+            char_count = len(source.text[node.idx:node.endidx])
+            print("SOURCE:")
+            print(source.text[node.idx:node.endidx])
+            rename_tokens(ctxt, node, True)
+            minified_source = minify_code(source.text[node.idx:node.endidx], ctxt, node, minify=True)
+            min_char_count = len(minified_source)
+            prefix_newline = "\n" if node.idx > 0 and source.text[node.idx-1] != "\n" else ""
+            source.text = f'{source.text[:node.idx]}{prefix_newline}-- T:{token_count} C:{char_count} minC:{min_char_count}\n-- minified source:\n{indent(minified_source,"-- ")}\n{source.text[node.idx:]}'
+    apply_node_tree(root, comment_before_function)
+
+from pico_tokenize import tokenize, count_tokens
+from pico_parse import NodeType
+from pico_process import PicoSource
+from pico_minify import minify_code
+from pico_rename import rename_tokens

pico_process.py

@@ -143,16 +143,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, unminify=False, fail=True, want_count=True):
+def process_code(ctxt, source, input_count=False, count=False, lint=False, minify=False, rename=False, unminify=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_unminify = unminify not in (None, False)
+    need_annotate = annotate not in (None, False)
 
-    if not need_lint and not need_minify and not need_unminify and not (want_count and (count or input_count)):
+    need_parse = need_lint or need_minify or need_unminify or need_annotate
+
+    if not need_parse and not (want_count and (count or input_count)):
         return True, ()
-    
-    need_parse = need_lint or need_minify or need_unminify
+
     need_all_comments = need_unminify or (need_minify and minify_needs_comments(minify))
 
     ok = False
@@ -182,6 +184,9 @@ def process_code(ctxt, source, input_count=False, count=False, lint=False, minif
             new_tokens = root.get_tokens() if need_parse else tokens
             print_token_count(count_tokens(new_tokens), handler=count)
 
+        if annotate:
+            annotate_code(ctxt, source, root)
+
     if fail and errors:
         throw("\n".join(map(str, errors)))
     return ok, errors
@@ -199,6 +204,7 @@ def echo_code(code, echo=True):
 from pico_lint import lint_code
 from pico_minify import minify_code, minify_needs_comments
 from pico_unminify import unminify_code
+from pico_annotate import annotate_code
 from pico_rename import rename_tokens
 
 # re-export some things for examples/etc.

shrinko8.py

@@ -52,6 +52,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=())
@@ -195,7 +198,8 @@ def main_inner(raw_args):
     ok, errors = process_code(ctxt, src, input_count=args.input_count, count=args.count,
                               lint=args.lint, minify=args.minify, rename=args.rename,
                               unminify=args.unminify,
-                              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: