Skip to content

Commit 9a84105

Browse files
authored
feat: refactor + improve stability (#1585)
## TLDR Problem <!-- Who is this for and what problem does it solve? --> <!-- Closes #ISSUE_ID --> ## Changes <!-- Describe what you tested -- manual steps, automated tests, or both. --> <!-- If you're an agent, only list tests you actually ran. -->
1 parent 35b2c15 commit 9a84105

15 files changed

Lines changed: 1342 additions & 49 deletions

packages/enricher/README.md

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
# @posthog/enricher
2+
3+
Detect and enrich PostHog SDK usage in source code. Uses tree-sitter AST analysis to find `capture()` calls, feature flag checks, `init()` calls, and variant branches across JavaScript, TypeScript, Python, Go, and Ruby.
4+
5+
## Quick start
6+
7+
```typescript
8+
import { PostHogEnricher } from "@posthog/enricher";
9+
10+
const enricher = new PostHogEnricher();
11+
await enricher.initialize("/path/to/grammars");
12+
13+
const result = await enricher.parse(sourceCode, "typescript");
14+
15+
result.events; // [{ name: "purchase", line: 5, dynamic: false }]
16+
result.flagChecks; // [{ method: "getFeatureFlag", flagKey: "new-checkout", line: 8 }]
17+
result.flagKeys; // ["new-checkout"]
18+
result.eventNames; // ["purchase"]
19+
result.toList(); // [{ type: "event", line: 5, name: "purchase", method: "capture" }, ...]
20+
```
21+
22+
## Enriching from the PostHog API
23+
24+
Let the enricher fetch everything it needs based on what `parse()` found — feature flags, experiments, event definitions, and event volume/user stats:
25+
26+
```typescript
27+
const result = await enricher.parse(sourceCode, "typescript");
28+
const enriched = await result.enrichFromApi({
29+
apiKey: "phx_...",
30+
host: "https://us.posthog.com",
31+
projectId: 12345,
32+
});
33+
34+
// Flags with staleness, rollout, experiment info
35+
enriched.enrichedFlags;
36+
// [{ flagKey: "new-checkout", flagType: "boolean", staleness: "fully_rolled_out",
37+
// rollout: 100, experiment: { name: "Checkout v2", ... }, ... }]
38+
39+
// Events with definition, volume, unique users
40+
enriched.enrichedEvents;
41+
// [{ eventName: "purchase", verified: true, lastSeenAt: "2025-04-01",
42+
// tags: ["revenue"], stats: { volume: 12500, uniqueUsers: 3200 }, ... }]
43+
44+
// Flat list combining both
45+
enriched.toList();
46+
// [{ type: "event", name: "purchase", verified: true, volume: 12500, ... },
47+
// { type: "flag", name: "new-checkout", flagType: "boolean", staleness: "fully_rolled_out", ... }]
48+
49+
// Source code with inline annotation comments
50+
enriched.toComments();
51+
// // [PostHog] Event: "purchase" (verified) — 12,500 events — 3,200 users
52+
// posthog.capture("purchase", { amount: 99 });
53+
//
54+
// // [PostHog] Flag: "new-checkout" — boolean — 100% rolled out — STALE (fully_rolled_out)
55+
// const flag = posthog.getFeatureFlag("new-checkout");
56+
```
57+
58+
## Supported languages
59+
60+
| Language | ID | Capture | Flags | Init | Variants |
61+
|---|---|---|---|---|---|
62+
| JavaScript | `javascript` | yes | yes | yes | yes |
63+
| TypeScript | `typescript` | yes | yes | yes | yes |
64+
| JSX | `javascriptreact` | yes | yes | yes | yes |
65+
| TSX | `typescriptreact` | yes | yes | yes | yes |
66+
| Python | `python` | yes | yes | yes | yes |
67+
| Go | `go` | yes | yes | yes | yes |
68+
| Ruby | `ruby` | yes | yes | yes | yes |
69+
70+
## API reference
71+
72+
### `PostHogEnricher`
73+
74+
Main entry point. Owns the tree-sitter parser lifecycle.
75+
76+
```typescript
77+
const enricher = new PostHogEnricher();
78+
await enricher.initialize(wasmDir);
79+
const result = await enricher.parse(source, languageId);
80+
enricher.dispose();
81+
```
82+
83+
### `ParseResult`
84+
85+
Returned by `enricher.parse()`. Contains all detected PostHog SDK usage.
86+
87+
| Property / Method | Type | Description |
88+
|---|---|---|
89+
| `calls` | `PostHogCall[]` | All detected SDK method calls |
90+
| `initCalls` | `PostHogInitCall[]` | `posthog.init()` and constructor calls |
91+
| `flagAssignments` | `FlagAssignment[]` | Flag result variable assignments |
92+
| `variantBranches` | `VariantBranch[]` | If/switch branches on flag values |
93+
| `functions` | `FunctionInfo[]` | Function definitions in the file |
94+
| `events` | `CapturedEvent[]` | Capture calls only |
95+
| `flagChecks` | `FlagCheck[]` | Flag method calls only |
96+
| `flagKeys` | `string[]` | Unique flag keys |
97+
| `eventNames` | `string[]` | Unique event names |
98+
| `toList()` | `ListItem[]` | Flat sorted list of all SDK usage |
99+
| `enrichFromApi(config)` | `Promise<EnrichedResult>` | Fetch from PostHog API and enrich |
100+
101+
### `EnrichedResult`
102+
103+
Returned by `enrich()` or `enrichFromApi()`. Detection combined with PostHog context.
104+
105+
| Property / Method | Type | Description |
106+
|---|---|---|
107+
| `enrichedFlags` | `EnrichedFlag[]` | Flags grouped by key with type, staleness, rollout, experiment |
108+
| `enrichedEvents` | `EnrichedEvent[]` | Events grouped by name with definition, stats, tags |
109+
| `toList()` | `EnrichedListItem[]` | Flat list with all metadata |
110+
| `toComments()` | `string` | Source code with inline annotation comments |
111+
112+
### `EnricherApiConfig`
113+
114+
```typescript
115+
interface EnricherApiConfig {
116+
apiKey: string;
117+
host: string; // e.g. "https://us.posthog.com"
118+
projectId: number;
119+
}
120+
```
121+
122+
### `EnrichedFlag`
123+
124+
```typescript
125+
interface EnrichedFlag {
126+
flagKey: string;
127+
flagType: "boolean" | "multivariate" | "remote_config";
128+
staleness: StalenessReason | null;
129+
rollout: number | null;
130+
variants: { key: string; rollout_percentage: number }[];
131+
flag: FeatureFlag | undefined;
132+
experiment: Experiment | undefined;
133+
occurrences: FlagCheck[];
134+
}
135+
```
136+
137+
### `EnrichedEvent`
138+
139+
```typescript
140+
interface EnrichedEvent {
141+
eventName: string;
142+
verified: boolean;
143+
lastSeenAt: string | null;
144+
tags: string[];
145+
stats: { volume?: number; uniqueUsers?: number } | undefined;
146+
definition: EventDefinition | undefined;
147+
occurrences: CapturedEvent[];
148+
}
149+
```
150+
151+
## Detection API
152+
153+
The lower-level detection API is also exported for direct use (this is the same API used by the PostHog VSCode extension):
154+
155+
```typescript
156+
import { PostHogDetector } from "@posthog/enricher";
157+
158+
const detector = new PostHogDetector();
159+
await detector.initialize(wasmDir);
160+
161+
const calls = await detector.findPostHogCalls(source, "typescript");
162+
const initCalls = await detector.findInitCalls(source, "typescript");
163+
const branches = await detector.findVariantBranches(source, "typescript");
164+
const assignments = await detector.findFlagAssignments(source, "typescript");
165+
const functions = await detector.findFunctions(source, "typescript");
166+
167+
detector.dispose();
168+
```
169+
170+
### Flag classification utilities
171+
172+
```typescript
173+
import { classifyFlagType, classifyStaleness } from "@posthog/enricher";
174+
175+
classifyFlagType(flag); // "boolean" | "multivariate" | "remote_config"
176+
classifyStaleness(key, flag, experiments, opts); // StalenessReason | null
177+
```
178+
179+
## Logging
180+
181+
Warnings are silenced by default. To receive them:
182+
183+
```typescript
184+
import { setLogger } from "@posthog/enricher";
185+
186+
setLogger({ warn: console.warn });
187+
```
188+
189+
## Setup
190+
191+
The package requires pre-built tree-sitter WASM grammar files. Run `pnpm fetch-grammars` to build them, or place pre-built `.wasm` files in the `grammars/` directory.

packages/enricher/package.json

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,18 +7,6 @@
77
".": {
88
"types": "./dist/index.d.ts",
99
"import": "./dist/index.js"
10-
},
11-
"./classification": {
12-
"types": "./dist/flag-classification.d.ts",
13-
"import": "./dist/flag-classification.js"
14-
},
15-
"./stale-flags": {
16-
"types": "./dist/stale-flags.d.ts",
17-
"import": "./dist/stale-flags.js"
18-
},
19-
"./types": {
20-
"types": "./dist/types.d.ts",
21-
"import": "./dist/types.js"
2210
}
2311
},
2412
"scripts": {
Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
import type { EnrichedEvent, EnrichedFlag, EnrichedListItem } from "./types.js";
2+
3+
function commentPrefix(languageId: string): string {
4+
if (languageId === "python" || languageId === "ruby") {
5+
return "#";
6+
}
7+
return "//";
8+
}
9+
10+
function formatFlagComment(flag: EnrichedFlag): string {
11+
const parts: string[] = [`Flag: "${flag.flagKey}"`];
12+
13+
if (flag.flag) {
14+
parts.push(flag.flagType);
15+
if (flag.rollout !== null) {
16+
parts.push(`${flag.rollout}% rolled out`);
17+
}
18+
if (flag.experiment) {
19+
const status = flag.experiment.end_date ? "complete" : "running";
20+
parts.push(`Experiment: "${flag.experiment.name}" (${status})`);
21+
}
22+
if (flag.staleness) {
23+
parts.push(`STALE (${flag.staleness})`);
24+
}
25+
}
26+
27+
return parts.join(" \u2014 ");
28+
}
29+
30+
function formatEventComment(event: EnrichedEvent): string {
31+
const parts: string[] = [`Event: "${event.eventName}"`];
32+
if (event.verified) {
33+
parts.push("(verified)");
34+
}
35+
if (event.stats?.volume !== undefined) {
36+
parts.push(`${event.stats.volume.toLocaleString()} events`);
37+
}
38+
if (event.stats?.uniqueUsers !== undefined) {
39+
parts.push(`${event.stats.uniqueUsers.toLocaleString()} users`);
40+
}
41+
if (event.definition?.description) {
42+
parts.push(event.definition.description);
43+
}
44+
return parts.join(" \u2014 ");
45+
}
46+
47+
export function formatComments(
48+
source: string,
49+
languageId: string,
50+
items: EnrichedListItem[],
51+
enrichedFlags: Map<string, EnrichedFlag>,
52+
enrichedEvents: Map<string, EnrichedEvent>,
53+
): string {
54+
const prefix = commentPrefix(languageId);
55+
const lines = source.split("\n");
56+
const sorted = [...items].sort((a, b) => a.line - b.line);
57+
58+
let offset = 0;
59+
// One comment per original source line — if multiple detections share a line,
60+
// only the first (by sort order) gets an annotation to keep output readable.
61+
const annotatedLines = new Set<number>();
62+
63+
for (const item of sorted) {
64+
const targetLine = item.line + offset;
65+
if (annotatedLines.has(item.line)) {
66+
continue;
67+
}
68+
annotatedLines.add(item.line);
69+
70+
let comment: string | null = null;
71+
72+
if (item.type === "flag") {
73+
const flag = enrichedFlags.get(item.name);
74+
if (flag) {
75+
comment = `${prefix} [PostHog] ${formatFlagComment(flag)}`;
76+
}
77+
} else if (item.type === "event") {
78+
const event = enrichedEvents.get(item.name);
79+
if (event) {
80+
comment = `${prefix} [PostHog] ${formatEventComment(event)}`;
81+
} else if (item.detail) {
82+
comment = `${prefix} [PostHog] Event: ${item.detail}`;
83+
}
84+
} else if (item.type === "init") {
85+
comment = `${prefix} [PostHog] Init: token "${item.name}"`;
86+
}
87+
88+
if (comment) {
89+
const indent = lines[targetLine]?.match(/^(\s*)/)?.[1] ?? "";
90+
lines.splice(targetLine, 0, `${indent}${comment}`);
91+
offset++;
92+
}
93+
}
94+
95+
return lines.join("\n");
96+
}

0 commit comments

Comments
 (0)