Skills › BMAD Method › bmad-toolsmith › count_tokens.py
count_tokens.py
103 lines · skills/bmad-toolsmith/scripts/count_tokens.py at bd70c64
1#!/usr/bin/env python3
2# /// script
3# requires-python = ">=3.11"
4# dependencies = ["tiktoken"]
5# ///
6"""count_tokens: the one length metric for skill authoring.
8Counts tokens with tiktoken's cl100k_base encoding. When tiktoken is not
9installed it falls back to len(text) // 4 and says so, so the script runs
10under a bare interpreter too.
12Usage:
13 count_tokens.py <path> [<path> ...] files, or directories recursed for text files
14 count_tokens.py --stdin count the text read from stdin
16Output, one JSON object on stdout:
17 {"files": [{"path": str, "tokens": int}, ...], "total": int, "method": "tiktoken" | "fallback"}
18"""
20from __future__ import annotations
22import argparse
23import json
24import sys
25from pathlib import Path
27ENCODING = "cl100k_base"
28TEXT_SUFFIXES = {".md", ".py", ".toml", ".yaml", ".yml", ".json", ".txt", ".csv", ".html", ".sh", ".cfg", ".ini"}
29SKIP_DIRS = {".git", "__pycache__", ".pytest_cache", "node_modules", ".venv", "venv"}
32def count_tokens(text: str) -> tuple[int, str]:
33 """Return (token_count, method); method is "fallback" when tiktoken is unavailable."""
34 try:
35 import tiktoken
36 except Exception:
37 return len(text) // 4, "fallback"
38 try:
39 enc = tiktoken.get_encoding(ENCODING)
40 except Exception:
41 return len(text) // 4, "fallback"
42 return len(enc.encode(text)), "tiktoken"
45def iter_text_files(root: Path):
46 """Yield text files under root, skipping noise directories and hidden files."""
47 # Sorted by the posix string so the order is the same on every OS; Path sorts case-insensitively on Windows.
48 for path in sorted(root.rglob("*"), key=lambda p: p.relative_to(root).as_posix()):
49 if not path.is_file():
50 continue
51 parts = path.relative_to(root).parts
52 if any(part in SKIP_DIRS or part.startswith(".") for part in parts):
53 continue
54 if path.suffix.lower() in TEXT_SUFFIXES:
55 yield path
58def read_text(path: Path) -> str:
59 try:
60 return path.read_text(encoding="utf-8")
61 except UnicodeDecodeError:
62 return path.read_text(encoding="utf-8", errors="replace")
65def count_paths(paths: list[Path]) -> dict:
66 files = []
67 method = "tiktoken"
68 for given in paths:
69 targets = list(iter_text_files(given)) if given.is_dir() else [given]
70 for target in targets:
71 tokens, used = count_tokens(read_text(target))
72 if used == "fallback":
73 method = "fallback"
74 files.append({"path": str(target), "tokens": tokens})
75 return {"files": files, "total": sum(f["tokens"] for f in files), "method": method}
78def main(argv: list[str] | None = None) -> int:
79 p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
80 p.add_argument("paths", nargs="*", type=Path, help="files or directories to count")
81 p.add_argument("--stdin", action="store_true", help="read text from stdin instead of paths")
82 args = p.parse_args(argv)
84 if args.stdin and args.paths:
85 p.error("provide paths or --stdin, not both")
86 if not args.stdin and not args.paths:
87 p.error("provide at least one path or --stdin")
89 if args.stdin:
90 tokens, method = count_tokens(sys.stdin.read())
91 result = {"files": [{"path": "<stdin>", "tokens": tokens}], "total": tokens, "method": method}
92 else:
93 missing = [str(path) for path in args.paths if not path.exists()]
94 if missing:
95 p.error(f"not found: {', '.join(missing)}")
96 result = count_paths(args.paths)
98 print(json.dumps(result))
99 return 0
102if __name__ == "__main__":
103 sys.exit(main())