@kt3k/tku 1.0.8 → 1.0.10

package/dist/main.js

@@ -60,7 +60,9 @@ async function listTextFiles(repoPath, options = {}) {
   const cwd = resolve(repoPath);
   const files = await listFiles(repoPath, options);
   const results = [];
-  for (const file of files) {
+  for (let i = 0; i < files.length; i++) {
+    const file = files[i];
+    options.onProgress?.(file, i + 1, files.length);
     const fullPath = resolve(cwd, file);
     try {
       if (!await isBinary(fullPath)) {
@@ -73,33 +75,61 @@ async function listTextFiles(repoPath, options = {}) {
 }
 
 // src/tokenize.ts
-import { readFile } from "node:fs/promises";
 import { resolve as resolve2 } from "node:path";
+import { availableParallelism } from "node:os";
+import { Worker } from "node:worker_threads";
 import { get_encoding } from "tiktoken";
-async function tokenizeFiles(repoPath, files, encoding) {
-  const enc = get_encoding(encoding);
-  try {
-    const results = [];
-    let totalTokens = 0;
-    for (const file of files) {
-      const fullPath = resolve2(repoPath, file);
-      const content = await readFile(fullPath, "utf-8");
-      const tokens = enc.encode_ordinary(content).length;
-      results.push({
-        path: file,
-        tokens
-      });
-      totalTokens += tokens;
-    }
-    return {
-      encoding,
-      files: results,
-      totalTokens,
-      totalFiles: results.length
-    };
-  } finally {
-    enc.free();
+var workerExt = import.meta.url.endsWith(".ts") ? ".ts" : ".js";
+var workerUrl = new URL(`./tokenize-worker${workerExt}`, import.meta.url);
+async function tokenizeFiles(repoPath, files, encoding, options = {}) {
+  const resolvedRepoPath = resolve2(repoPath);
+  const numWorkers = Math.min(availableParallelism(), files.length || 1);
+  const chunks = Array.from({
+    length: numWorkers
+  }, () => []);
+  for (let i = 0; i < files.length; i++) {
+    chunks[i % numWorkers].push(files[i]);
   }
+  let progressCount = 0;
+  const total = files.length;
+  const workerPromises = chunks.map((chunk) => {
+    if (chunk.length === 0) return Promise.resolve([]);
+    return new Promise((resolvePromise, reject) => {
+      const worker = new Worker(workerUrl);
+      worker.on("message", (msg) => {
+        if (msg.type === "progress") {
+          progressCount++;
+          options.onProgress?.(msg.file, progressCount, total);
+        } else if (msg.type === "done") {
+          resolvePromise(msg.results);
+          worker.terminate();
+        }
+      });
+      worker.on("error", (err) => {
+        reject(err);
+        worker.terminate();
+      });
+      worker.postMessage({
+        repoPath: resolvedRepoPath,
+        files: chunk,
+        encoding
+      });
+    });
+  });
+  const chunkResults = await Promise.all(workerPromises);
+  const results = chunkResults.flat();
+  const orderMap = new Map(files.map((f, i) => [
+    f,
+    i
+  ]));
+  results.sort((a, b) => orderMap.get(a.path) - orderMap.get(b.path));
+  const totalTokens = results.reduce((sum, f) => sum + f.tokens, 0);
+  return {
+    encoding,
+    files: results,
+    totalTokens,
+    totalFiles: results.length
+  };
 }
 
 // src/format.ts
@@ -112,7 +142,7 @@ function formatTokenCount(n) {
   }
   return String(n);
 }
-function formatTable(result) {
+function formatTable(result, omitted = 0) {
   const lines = [];
   const formatted = result.files.map((f) => ({
     path: f.path,
@@ -123,6 +153,9 @@ function formatTable(result) {
   for (const f of formatted) {
     lines.push(`${f.display.padStart(maxWidth)}  ${f.path}`);
   }
+  if (omitted > 0) {
+    lines.push(`${"".padStart(maxWidth)}  ... ${omitted} more files`);
+  }
   lines.push(`${"\u2500".repeat(maxWidth)}\u2500\u2500`);
   const totalDisplay = formatTokenCount(result.totalTokens);
   lines.push(`${totalDisplay.padStart(maxWidth)}  total (${result.totalFiles} files)`);
@@ -139,11 +172,12 @@ function formatResult(result, options = {}) {
     sorted.sort((a, b) => b.tokens - a.tokens);
   }
   const filtered = top !== void 0 ? sorted.slice(0, top) : sorted;
+  const omitted = sorted.length - filtered.length;
   const adjusted = {
     ...result,
     files: filtered
   };
-  return json ? JSON.stringify(adjusted, null, 2) : formatTable(adjusted);
+  return json ? JSON.stringify(adjusted, null, 2) : formatTable(adjusted, omitted);
 }
 
 // src/main.ts
@@ -203,12 +237,28 @@ async function main() {
   const encoding = values.encoding;
   const sort = values.sort;
   const top = values.top !== void 0 ? Number(values.top) : void 0;
+  const isTTY = process.stderr.isTTY;
+  function status(msg) {
+    if (isTTY) {
+      process.stderr.write(`\r\x1B[K${msg}`);
+    }
+  }
+  function clearStatus() {
+    if (isTTY) {
+      process.stderr.write("\r\x1B[K");
+    }
+  }
   try {
     const files = await listTextFiles(repoPath, {
       exclude: values.exclude,
-      noGitignore: !values.gitignore
+      noGitignore: !values.gitignore,
+      onProgress: (file, i, total) => status(`Scanning [${i}/${total}] ${file}`)
+    });
+    clearStatus();
+    const result = await tokenizeFiles(repoPath, files, encoding, {
+      onProgress: (file, i, total) => status(`Tokenizing [${i}/${total}] ${file}`)
     });
-    const result = await tokenizeFiles(repoPath, files, encoding);
+    clearStatus();
     const output = formatResult(result, {
       json: values.json,
       top,

package/dist/tokenize-worker.js

@@ -0,0 +1,30 @@
+// src/tokenize-worker.ts
+import { parentPort } from "node:worker_threads";
+import { readFile } from "node:fs/promises";
+import { resolve } from "node:path";
+import { get_encoding } from "tiktoken";
+parentPort.on("message", async (msg) => {
+  const enc = get_encoding(msg.encoding);
+  try {
+    const results = [];
+    for (const file of msg.files) {
+      const fullPath = resolve(msg.repoPath, file);
+      const content = await readFile(fullPath, "utf-8");
+      const tokens = enc.encode_ordinary(content).length;
+      results.push({
+        path: file,
+        tokens
+      });
+      parentPort.postMessage({
+        type: "progress",
+        file
+      });
+    }
+    parentPort.postMessage({
+      type: "done",
+      results
+    });
+  } finally {
+    enc.free();
+  }
+});

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@kt3k/tku",
-  "version": "1.0.8",
+  "version": "1.0.10",
   "description": "",
   "main": "dist/main.js",
   "bin": {
@@ -10,8 +10,8 @@
     "dist"
   ],
   "scripts": {
-    "build": "deno bundle --external tiktoken --external picomatch -o dist/main.js src/main.ts && chmod +x dist/main.js",
-    "prepublish": "npm run build",
+    "build": "deno bundle --external tiktoken --external picomatch -o dist/main.js src/main.ts && deno bundle --external tiktoken -o dist/tokenize-worker.js src/tokenize-worker.ts && chmod +x dist/main.js",
+    "prepublishOnly": "npm run build",
     "test": "vitest run"
   },
   "keywords": [],