Skip to content

Commit c8faf81

Browse files
cursoragentcursoragent
committed
Fix type annotation parsing and update test fixtures for Python tracking
1 parent 4fb4ec4 commit c8faf81

File tree

3 files changed

+34
-19
lines changed

3 files changed

+34
-19
lines changed

‎src/analyze/python/pythonTrackingAnalyzer.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -164,8 +164,11 @@ def extract_type_annotation(self, annotation: ast.AST) -> PropertyType:
164164

165165
elif isinstance(annotation, ast.Subscript):
166166
# Handle generic types like List[int], Dict[str, int]
167-
if hasattr(annotation.value, 'id'):
168-
container_type = annotation.value.id
167+
container_type = getattr(annotation.value, 'id', None)
168+
if container_type is None and isinstance(annotation.value, ast.Attribute):
169+
container_type = getattr(annotation.value.attr, 'id', None) or annotation.value.attr
170+
171+
if container_type:
169172

170173
if container_type in ARRAY_TYPES:
171174
# Try to get the type parameter for arrays

‎tests/analyzePython.test.js‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ test.describe('analyzePythonFile', () => {
2222
assert.ok(segmentEvent);
2323
assert.strictEqual(segmentEvent.source, 'segment');
2424
assert.strictEqual(segmentEvent.functionName, 'segment_track');
25-
assert.strictEqual(segmentEvent.line, 11);
25+
assert.strictEqual(segmentEvent.line, 23);
2626
assert.deepStrictEqual(segmentEvent.properties, {
2727
user_id: { type: 'string' },
2828
method: { type: 'string' },
@@ -35,7 +35,7 @@ test.describe('analyzePythonFile', () => {
3535
assert.ok(mixpanelEvent);
3636
assert.strictEqual(mixpanelEvent.source, 'mixpanel');
3737
assert.strictEqual(mixpanelEvent.functionName, 'mixpanel_track');
38-
assert.strictEqual(mixpanelEvent.line, 21);
38+
assert.strictEqual(mixpanelEvent.line, 33);
3939
assert.deepStrictEqual(mixpanelEvent.properties, {
4040
distinct_id: { type: 'string' },
4141
plan: { type: 'string' },
@@ -48,7 +48,7 @@ test.describe('analyzePythonFile', () => {
4848
assert.ok(amplitudeEvent);
4949
assert.strictEqual(amplitudeEvent.source, 'amplitude');
5050
assert.strictEqual(amplitudeEvent.functionName, 'amplitude_track');
51-
assert.strictEqual(amplitudeEvent.line, 31);
51+
assert.strictEqual(amplitudeEvent.line, 43);
5252
assert.deepStrictEqual(amplitudeEvent.properties, {
5353
user_id: { type: 'string' },
5454
color: { type: 'string' },
@@ -60,7 +60,7 @@ test.describe('analyzePythonFile', () => {
6060
assert.ok(rudderstackEvent);
6161
assert.strictEqual(rudderstackEvent.source, 'rudderstack');
6262
assert.strictEqual(rudderstackEvent.functionName, 'rudderstack_track');
63-
assert.strictEqual(rudderstackEvent.line, 47);
63+
assert.strictEqual(rudderstackEvent.line, 59);
6464
assert.deepStrictEqual(rudderstackEvent.properties, {
6565
user_id: { type: 'string' },
6666
timestamp: { type: 'number' },
@@ -73,7 +73,7 @@ test.describe('analyzePythonFile', () => {
7373
assert.ok(posthogEvent1);
7474
assert.strictEqual(posthogEvent1.source, 'posthog');
7575
assert.strictEqual(posthogEvent1.functionName, 'posthog_capture');
76-
assert.strictEqual(posthogEvent1.line, 58);
76+
assert.strictEqual(posthogEvent1.line, 70);
7777
assert.deepStrictEqual(posthogEvent1.properties, {
7878
method: { type: 'string' },
7979
is_free_trial: { type: 'boolean' },
@@ -84,7 +84,7 @@ test.describe('analyzePythonFile', () => {
8484
assert.ok(posthogEvent2);
8585
assert.strictEqual(posthogEvent2.source, 'posthog');
8686
assert.strictEqual(posthogEvent2.functionName, 'posthog_capture');
87-
assert.strictEqual(posthogEvent2.line, 64);
87+
assert.strictEqual(posthogEvent2.line, 76);
8888
assert.deepStrictEqual(posthogEvent2.properties, {
8989
method: { type: 'string' },
9090
is_free_trial: { type: 'boolean' },
@@ -96,7 +96,7 @@ test.describe('analyzePythonFile', () => {
9696
assert.ok(snowplowEvent);
9797
assert.strictEqual(snowplowEvent.source, 'snowplow');
9898
assert.strictEqual(snowplowEvent.functionName, 'snowplow_track_events');
99-
assert.strictEqual(snowplowEvent.line, 74);
99+
assert.strictEqual(snowplowEvent.line, 86);
100100
assert.deepStrictEqual(snowplowEvent.properties, {
101101
category: { type: 'string' },
102102
label: { type: 'string' },
@@ -109,7 +109,7 @@ test.describe('analyzePythonFile', () => {
109109
assert.ok(customEvent);
110110
assert.strictEqual(customEvent.source, 'custom');
111111
assert.strictEqual(customEvent.functionName, 'main');
112-
assert.strictEqual(customEvent.line, 89);
112+
assert.strictEqual(customEvent.line, 101);
113113
assert.deepStrictEqual(customEvent.properties, {
114114
userId: { type: 'string' },
115115
key: { type: 'string' },
@@ -178,7 +178,7 @@ test.describe('analyzePythonFile', () => {
178178
assert.ok(eventMap['User Signed Up']);
179179
assert.strictEqual(eventMap['User Signed Up'].eventName, 'User Signed Up');
180180
assert.strictEqual(eventMap['User Signed Up'].source, 'segment');
181-
assert.strictEqual(eventMap['User Signed Up'].line, 11);
181+
assert.strictEqual(eventMap['User Signed Up'].line, 23);
182182
assert.strictEqual(eventMap['User Signed Up'].functionName, 'segment_track');
183183
assert.deepStrictEqual(eventMap['User Signed Up'].properties, {
184184
user_id: { type: 'string' },
@@ -190,7 +190,7 @@ test.describe('analyzePythonFile', () => {
190190
assert.ok(eventMap['Purchase Completed']);
191191
assert.strictEqual(eventMap['Purchase Completed'].eventName, 'Purchase Completed');
192192
assert.strictEqual(eventMap['Purchase Completed'].source, 'mixpanel');
193-
assert.strictEqual(eventMap['Purchase Completed'].line, 21);
193+
assert.strictEqual(eventMap['Purchase Completed'].line, 33);
194194
assert.strictEqual(eventMap['Purchase Completed'].functionName, 'mixpanel_track');
195195
assert.deepStrictEqual(eventMap['Purchase Completed'].properties, {
196196
distinct_id: { type: 'string' },
@@ -202,7 +202,7 @@ test.describe('analyzePythonFile', () => {
202202
assert.ok(eventMap['Button Clicked']);
203203
assert.strictEqual(eventMap['Button Clicked'].eventName, 'Button Clicked');
204204
assert.strictEqual(eventMap['Button Clicked'].source, 'amplitude');
205-
assert.strictEqual(eventMap['Button Clicked'].line, 31);
205+
assert.strictEqual(eventMap['Button Clicked'].line, 43);
206206
assert.strictEqual(eventMap['Button Clicked'].functionName, 'amplitude_track');
207207
assert.deepStrictEqual(eventMap['Button Clicked'].properties, {
208208
user_id: { type: 'string' },
@@ -213,7 +213,7 @@ test.describe('analyzePythonFile', () => {
213213
assert.ok(eventMap['User Logged In']);
214214
assert.strictEqual(eventMap['User Logged In'].eventName, 'User Logged In');
215215
assert.strictEqual(eventMap['User Logged In'].source, 'rudderstack');
216-
assert.strictEqual(eventMap['User Logged In'].line, 47);
216+
assert.strictEqual(eventMap['User Logged In'].line, 59);
217217
assert.strictEqual(eventMap['User Logged In'].functionName, 'rudderstack_track');
218218
assert.deepStrictEqual(eventMap['User Logged In'].properties, {
219219
user_id: { type: 'string' },
@@ -225,7 +225,7 @@ test.describe('analyzePythonFile', () => {
225225
assert.ok(eventMap['user_signed_up']);
226226
assert.strictEqual(eventMap['user_signed_up'].eventName, 'user_signed_up');
227227
assert.strictEqual(eventMap['user_signed_up'].source, 'posthog');
228-
assert.strictEqual(eventMap['user_signed_up'].line, 58);
228+
assert.strictEqual(eventMap['user_signed_up'].line, 70);
229229
assert.strictEqual(eventMap['user_signed_up'].functionName, 'posthog_capture');
230230
assert.deepStrictEqual(eventMap['user_signed_up'].properties, {
231231
method: { type: 'string' },
@@ -236,7 +236,7 @@ test.describe('analyzePythonFile', () => {
236236
assert.ok(eventMap['user_cancelled_subscription']);
237237
assert.strictEqual(eventMap['user_cancelled_subscription'].eventName, 'user_cancelled_subscription');
238238
assert.strictEqual(eventMap['user_cancelled_subscription'].source, 'posthog');
239-
assert.strictEqual(eventMap['user_cancelled_subscription'].line, 64);
239+
assert.strictEqual(eventMap['user_cancelled_subscription'].line, 76);
240240
assert.strictEqual(eventMap['user_cancelled_subscription'].functionName, 'posthog_capture');
241241
assert.deepStrictEqual(eventMap['user_cancelled_subscription'].properties, {
242242
method: { type: 'string' },
@@ -247,7 +247,7 @@ test.describe('analyzePythonFile', () => {
247247
assert.ok(eventMap['add-to-basket']);
248248
assert.strictEqual(eventMap['add-to-basket'].eventName, 'add-to-basket');
249249
assert.strictEqual(eventMap['add-to-basket'].source, 'snowplow');
250-
assert.strictEqual(eventMap['add-to-basket'].line, 74);
250+
assert.strictEqual(eventMap['add-to-basket'].line, 86);
251251
assert.strictEqual(eventMap['add-to-basket'].functionName, 'snowplow_track_events');
252252
assert.deepStrictEqual(eventMap['add-to-basket'].properties, {
253253
category: { type: 'string' },
@@ -259,7 +259,7 @@ test.describe('analyzePythonFile', () => {
259259
assert.ok(eventMap['custom_event']);
260260
assert.strictEqual(eventMap['custom_event'].eventName, 'custom_event');
261261
assert.strictEqual(eventMap['custom_event'].source, 'custom');
262-
assert.strictEqual(eventMap['custom_event'].line, 89);
262+
assert.strictEqual(eventMap['custom_event'].line, 101);
263263
assert.strictEqual(eventMap['custom_event'].functionName, 'main');
264264
assert.deepStrictEqual(eventMap['custom_event'].properties, {
265265
userId: { type: 'string' },

‎tests/fixtures/python/main.py‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,17 @@
11
from typing import Any, Dict, List
22

3+
# Stub imports for external analytics SDKs so linters/type checkers don't complain.
4+
# They are not actually executed in tests.
5+
try:
6+
import segment # type: ignore
7+
import mixpanel # type: ignore
8+
import amplitude # type: ignore
9+
import posthog # type: ignore
10+
import rudderstack # type: ignore
11+
import snowplow_tracker # type: ignore
12+
except ImportError: # pragma: no cover
13+
pass
14+
315
# Custom tracking function stub
416
def customTrackFunction(user_id: str, event_name: str, params: Dict[str, Any]) -> None:
517
print(f"Custom track: {user_id} - {event_name} - {params}")
@@ -82,7 +94,7 @@ def snowplow_track_events(category: str, value: float) -> None:
8294
def main() -> None:
8395
segment_track("user123", plan="Pro")
8496
mixpanel_track("user123", 9.99, ["apple", "banana"])
85-
amplitude_track("ButtonClicked", {"color": "red", "size": 12})
97+
amplitude_track("user123", 12)
8698
rudderstack_track("user123", "iOS", 14)
8799
posthog_capture("user123", "email", True, "premium")
88100
snowplow_track_events("shop", 2)

0 commit comments

Comments
 (0)