ucn 3.4.0 → 3.4.2

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/core/project.js CHANGED
@@ -1009,11 +1009,33 @@ class ProjectIndex {
1009
1009
 
1010
1010
  // Smart method call handling
1011
1011
  if (call.isMethod) {
1012
- // Always skip this/self/cls calls (internal state access, not function calls)
1013
- if (['this', 'self', 'cls'].includes(call.receiver)) continue;
1014
- // Go doesn't use this/self/cls - always include Go method calls
1015
- // For other languages, skip method calls unless explicitly requested
1016
- if (fileEntry.language !== 'go' && !options.includeMethods) continue;
1012
+ if (call.selfAttribute && fileEntry.language === 'python') {
1013
+ // self.attr.method() — resolve via attribute type inference
1014
+ const callerSymbol = this.findEnclosingFunction(filePath, call.line, true);
1015
+ if (!callerSymbol?.className) continue;
1016
+ const attrTypes = this.getInstanceAttributeTypes(filePath, callerSymbol.className);
1017
+ if (!attrTypes) continue;
1018
+ const targetClass = attrTypes.get(call.selfAttribute);
1019
+ if (!targetClass) continue;
1020
+ // Check if any definition of searched function belongs to targetClass
1021
+ const matchesDef = definitions.some(d => d.className === targetClass);
1022
+ if (!matchesDef) continue;
1023
+ // Falls through to add as caller
1024
+ } else if (fileEntry.language === 'python' && ['self', 'cls'].includes(call.receiver)) {
1025
+ // self.method() / cls.method() — resolve to same-class method
1026
+ const callerSymbol = this.findEnclosingFunction(filePath, call.line, true);
1027
+ if (!callerSymbol?.className) continue;
1028
+ // Check if any definition of searched function belongs to caller's class
1029
+ const matchesDef = definitions.some(d => d.className === callerSymbol.className);
1030
+ if (!matchesDef) continue;
1031
+ // Falls through to add as caller
1032
+ } else {
1033
+ // Always skip this/self/cls calls (internal state access, not function calls)
1034
+ if (['this', 'self', 'cls'].includes(call.receiver)) continue;
1035
+ // Go doesn't use this/self/cls - always include Go method calls
1036
+ // For other languages, skip method calls unless explicitly requested
1037
+ if (fileEntry.language !== 'go' && !options.includeMethods) continue;
1038
+ }
1017
1039
  }
1018
1040
 
1019
1041
  // Skip definition lines
@@ -1110,6 +1132,8 @@ class ProjectIndex {
1110
1132
  const language = fileEntry?.language;
1111
1133
 
1112
1134
  const callees = new Map(); // key -> { name, bindingId, count }
1135
+ let selfAttrCalls = null; // collected for Python self.attr.method() resolution
1136
+ let selfMethodCalls = null; // collected for Python self.method() resolution
1113
1137
 
1114
1138
  for (const call of calls) {
1115
1139
  // Filter to calls within this function's scope using enclosingFunction
@@ -1119,14 +1143,36 @@ class ProjectIndex {
1119
1143
 
1120
1144
  // Smart method call handling:
1121
1145
  // - Go: include all method calls (Go doesn't use this/self/cls)
1146
+ // - Python self.method(): resolve to same-class method (handled below)
1147
+ // - Python self.attr.method(): resolve via selfAttribute (handled below)
1122
1148
  // - Other languages: skip method calls unless explicitly requested
1123
1149
  if (call.isMethod) {
1124
- if (language !== 'go' && !options.includeMethods) continue;
1150
+ if (call.selfAttribute && language === 'python') {
1151
+ // Will be resolved in second pass below
1152
+ } else if (language === 'python' && ['self', 'cls'].includes(call.receiver)) {
1153
+ // self.method() / cls.method() — resolve to same-class method below
1154
+ } else if (language !== 'go' && !options.includeMethods) {
1155
+ continue;
1156
+ }
1125
1157
  }
1126
1158
 
1127
1159
  // Skip keywords and built-ins
1128
1160
  if (this.isKeyword(call.name, language)) continue;
1129
1161
 
1162
+ // Collect selfAttribute calls for second-pass resolution
1163
+ if (call.selfAttribute && language === 'python') {
1164
+ if (!selfAttrCalls) selfAttrCalls = [];
1165
+ selfAttrCalls.push(call);
1166
+ continue;
1167
+ }
1168
+
1169
+ // Collect Python self.method() calls for same-class resolution
1170
+ if (language === 'python' && call.isMethod && ['self', 'cls'].includes(call.receiver)) {
1171
+ if (!selfMethodCalls) selfMethodCalls = [];
1172
+ selfMethodCalls.push(call);
1173
+ continue;
1174
+ }
1175
+
1130
1176
  // Resolve binding within this file (without mutating cached call objects)
1131
1177
  let calleeKey = call.bindingId || call.name;
1132
1178
  let bindingResolved = call.bindingId;
@@ -1179,6 +1225,60 @@ class ProjectIndex {
1179
1225
  }
1180
1226
  }
1181
1227
 
1228
+ // Second pass: resolve Python self.attr.method() calls
1229
+ if (selfAttrCalls && def.className) {
1230
+ const attrTypes = this.getInstanceAttributeTypes(def.file, def.className);
1231
+ if (attrTypes) {
1232
+ for (const call of selfAttrCalls) {
1233
+ const targetClass = attrTypes.get(call.selfAttribute);
1234
+ if (!targetClass) continue;
1235
+
1236
+ // Find method in symbol table where className matches
1237
+ const symbols = this.symbols.get(call.name);
1238
+ if (!symbols) continue;
1239
+
1240
+ const match = symbols.find(s => s.className === targetClass);
1241
+ if (!match) continue;
1242
+
1243
+ const key = match.bindingId || `${targetClass}.${call.name}`;
1244
+ const existing = callees.get(key);
1245
+ if (existing) {
1246
+ existing.count += 1;
1247
+ } else {
1248
+ callees.set(key, {
1249
+ name: call.name,
1250
+ bindingId: match.bindingId,
1251
+ count: 1
1252
+ });
1253
+ }
1254
+ }
1255
+ }
1256
+ }
1257
+
1258
+ // Third pass: resolve Python self.method() calls to same-class methods
1259
+ if (selfMethodCalls && def.className) {
1260
+ for (const call of selfMethodCalls) {
1261
+ const symbols = this.symbols.get(call.name);
1262
+ if (!symbols) continue;
1263
+
1264
+ // Find method in same class
1265
+ const match = symbols.find(s => s.className === def.className);
1266
+ if (!match) continue;
1267
+
1268
+ const key = match.bindingId || `${def.className}.${call.name}`;
1269
+ const existing = callees.get(key);
1270
+ if (existing) {
1271
+ existing.count += 1;
1272
+ } else {
1273
+ callees.set(key, {
1274
+ name: call.name,
1275
+ bindingId: match.bindingId,
1276
+ count: 1
1277
+ });
1278
+ }
1279
+ }
1280
+ }
1281
+
1182
1282
  // Look up each callee in the symbol table
1183
1283
  // For methods, prefer callees from: 1) same file, 2) same package, 3) same receiver type
1184
1284
  const result = [];
@@ -1490,8 +1590,9 @@ class ProjectIndex {
1490
1590
  const fileEntry = this.files.get(filePath);
1491
1591
  if (!fileEntry) return null;
1492
1592
 
1593
+ const nonCallableTypes = new Set(['class', 'struct', 'interface', 'type', 'state']);
1493
1594
  for (const symbol of fileEntry.symbols) {
1494
- if (symbol.type === 'function' &&
1595
+ if (!nonCallableTypes.has(symbol.type) &&
1495
1596
  symbol.startLine <= lineNum &&
1496
1597
  symbol.endLine >= lineNum) {
1497
1598
  if (returnSymbol) {
@@ -1503,6 +1604,35 @@ class ProjectIndex {
1503
1604
  return null;
1504
1605
  }
1505
1606
 
1607
+ /**
1608
+ * Get instance attribute types for a class in a file.
1609
+ * Returns Map<attrName, typeName> for a given className.
1610
+ * Caches results per file.
1611
+ */
1612
+ getInstanceAttributeTypes(filePath, className) {
1613
+ if (!this._attrTypeCache) this._attrTypeCache = new Map();
1614
+
1615
+ let fileCache = this._attrTypeCache.get(filePath);
1616
+ if (!fileCache) {
1617
+ const fileEntry = this.files.get(filePath);
1618
+ if (!fileEntry || fileEntry.language !== 'python') return null;
1619
+
1620
+ const langModule = getLanguageModule('python');
1621
+ if (!langModule?.findInstanceAttributeTypes) return null;
1622
+
1623
+ try {
1624
+ const content = fs.readFileSync(filePath, 'utf-8');
1625
+ const parser = getParser('python');
1626
+ fileCache = langModule.findInstanceAttributeTypes(content, parser);
1627
+ this._attrTypeCache.set(filePath, fileCache);
1628
+ } catch {
1629
+ return null;
1630
+ }
1631
+ }
1632
+
1633
+ return fileCache.get(className) || null;
1634
+ }
1635
+
1506
1636
  /**
1507
1637
  * Extract type names from a function definition
1508
1638
  */
@@ -400,9 +400,14 @@ function findCallsInCode(code, parser) {
400
400
  traverseTree(tree.rootNode, (node) => {
401
401
  // Track function entry
402
402
  if (isFunctionNode(node)) {
403
+ // Use decorated_definition start line if present, to match symbol index
404
+ let startLine = node.startPosition.row + 1;
405
+ if (node.parent && node.parent.type === 'decorated_definition') {
406
+ startLine = node.parent.startPosition.row + 1;
407
+ }
403
408
  functionStack.push({
404
409
  name: extractFunctionName(node),
405
- startLine: node.startPosition.row + 1,
410
+ startLine,
406
411
  endLine: node.endPosition.row + 1
407
412
  });
408
413
  }
@@ -425,16 +430,32 @@ function findCallsInCode(code, parser) {
425
430
  uncertain
426
431
  });
427
432
  } else if (funcNode.type === 'attribute') {
428
- // Method/attribute call: obj.foo()
433
+ // Method/attribute call: obj.foo() or self.attr.foo()
429
434
  const attrNode = funcNode.childForFieldName('attribute');
430
435
  const objNode = funcNode.childForFieldName('object');
431
436
 
432
437
  if (attrNode) {
438
+ let receiver = objNode?.type === 'identifier' ? objNode.text : undefined;
439
+ let selfAttribute = undefined;
440
+
441
+ // Detect self.X.method() pattern: objNode is attribute access on self/cls
442
+ if (objNode?.type === 'attribute') {
443
+ const innerObj = objNode.childForFieldName('object');
444
+ const innerAttr = objNode.childForFieldName('attribute');
445
+ if (innerObj?.type === 'identifier' &&
446
+ ['self', 'cls'].includes(innerObj.text) &&
447
+ innerAttr) {
448
+ selfAttribute = innerAttr.text;
449
+ receiver = innerObj.text;
450
+ }
451
+ }
452
+
433
453
  calls.push({
434
454
  name: attrNode.text,
435
455
  line: node.startPosition.row + 1,
436
456
  isMethod: true,
437
- receiver: objNode?.type === 'identifier' ? objNode.text : undefined,
457
+ receiver,
458
+ ...(selfAttribute && { selfAttribute }),
438
459
  enclosingFunction,
439
460
  uncertain
440
461
  });
@@ -699,6 +720,183 @@ function findUsagesInCode(code, name, parser) {
699
720
  return usages;
700
721
  }
701
722
 
723
+ /**
724
+ * Find instance attribute types from __init__ constructor assignments.
725
+ * Parses self.X = ClassName(...) patterns in __init__ methods.
726
+ * @param {string} code - Source code to analyze
727
+ * @param {object} parser - Tree-sitter parser instance
728
+ * @returns {Map<string, Map<string, string>>} className -> (attrName -> typeName)
729
+ */
730
+ function findInstanceAttributeTypes(code, parser) {
731
+ const tree = parseTree(parser, code);
732
+ const result = new Map(); // className -> Map(attrName -> typeName)
733
+
734
+ const PRIMITIVE_TYPES = new Set(['int', 'float', 'str', 'bool', 'bytes', 'list', 'dict', 'set', 'tuple', 'None', 'Any', 'object']);
735
+
736
+ traverseTree(tree.rootNode, (node) => {
737
+ if (node.type !== 'class_definition') return true;
738
+
739
+ const classNameNode = node.childForFieldName('name');
740
+ if (!classNameNode) return true;
741
+ const className = classNameNode.text;
742
+
743
+ const body = node.childForFieldName('body');
744
+ if (!body) return false;
745
+
746
+ const attrTypes = new Map();
747
+
748
+ // Check for @dataclass decorator — scan annotated class-level fields
749
+ const parentNode = node.parent;
750
+ if (parentNode?.type === 'decorated_definition') {
751
+ for (let d = 0; d < parentNode.childCount; d++) {
752
+ const dec = parentNode.child(d);
753
+ if (dec.type !== 'decorator') continue;
754
+ // Match @dataclass or @dataclasses.dataclass
755
+ const decText = dec.text;
756
+ if (decText.startsWith('@dataclass') || decText.includes('.dataclass')) {
757
+ // Scan class body for annotated fields: name: Type = ...
758
+ for (let i = 0; i < body.childCount; i++) {
759
+ const stmt = body.child(i);
760
+ if (stmt.type !== 'expression_statement') continue;
761
+ const assign = stmt.firstChild;
762
+ if (!assign || assign.type !== 'assignment') continue;
763
+
764
+ // Must have a type annotation
765
+ const typeNode = assign.childForFieldName('type');
766
+ if (!typeNode) continue;
767
+
768
+ // Extract type name from annotation
769
+ const typeIdent = typeNode.type === 'type' ? typeNode.firstChild : typeNode;
770
+ if (!typeIdent || typeIdent.type !== 'identifier') continue;
771
+ const typeName = typeIdent.text;
772
+
773
+ // Skip primitives and lowercase types
774
+ if (PRIMITIVE_TYPES.has(typeName)) continue;
775
+ if (typeName[0] < 'A' || typeName[0] > 'Z') continue;
776
+
777
+ // Field name from LHS
778
+ const lhs = assign.childForFieldName('left');
779
+ if (!lhs || lhs.type !== 'identifier') continue;
780
+ attrTypes.set(lhs.text, typeName);
781
+ }
782
+ break;
783
+ }
784
+ }
785
+ }
786
+
787
+ // Scan __init__ for self.X = ClassName(...) assignments
788
+ for (let i = 0; i < body.childCount; i++) {
789
+ let child = body.child(i);
790
+ // Handle decorated_definition wrapper
791
+ if (child.type === 'decorated_definition') {
792
+ for (let j = 0; j < child.childCount; j++) {
793
+ if (child.child(j).type === 'function_definition') {
794
+ child = child.child(j);
795
+ break;
796
+ }
797
+ }
798
+ }
799
+ if (child.type !== 'function_definition') continue;
800
+
801
+ const fnName = child.childForFieldName('name');
802
+ if (!fnName || fnName.text !== '__init__') continue;
803
+
804
+ // Found __init__, now scan for self.X = ClassName(...) assignments
805
+ const initBody = child.childForFieldName('body');
806
+ if (!initBody) continue;
807
+
808
+ traverseTree(initBody, (stmt) => {
809
+ if (stmt.type !== 'expression_statement') return true;
810
+
811
+ const assign = stmt.firstChild;
812
+ if (!assign || assign.type !== 'assignment') return true;
813
+
814
+ // LHS: self.X
815
+ const lhs = assign.childForFieldName('left');
816
+ if (!lhs || lhs.type !== 'attribute') return true;
817
+ const lhsObj = lhs.childForFieldName('object');
818
+ const lhsAttr = lhs.childForFieldName('attribute');
819
+ if (!lhsObj || lhsObj.text !== 'self' || !lhsAttr) return true;
820
+
821
+ const attrName = lhsAttr.text;
822
+
823
+ // RHS: ClassName(...) or param or ClassName(...)
824
+ const rhs = assign.childForFieldName('right');
825
+ if (!rhs) return true;
826
+
827
+ const typeName = extractConstructorName(rhs);
828
+ if (typeName) {
829
+ attrTypes.set(attrName, typeName);
830
+ }
831
+
832
+ return true;
833
+ });
834
+ }
835
+
836
+ if (attrTypes.size > 0) {
837
+ result.set(className, attrTypes);
838
+ }
839
+
840
+ return false; // don't descend into nested classes from traverseTree
841
+ });
842
+
843
+ return result;
844
+ }
845
+
846
+ /**
847
+ * Extract constructor class name from an expression node.
848
+ * Handles: ClassName(...), param or ClassName(...), (param or ClassName(...)),
849
+ * expr if cond else ClassName(...)
850
+ */
851
+ function extractConstructorName(node) {
852
+ if (!node) return null;
853
+
854
+ // Direct call: ClassName(...)
855
+ if (node.type === 'call') {
856
+ const func = node.childForFieldName('function');
857
+ if (func?.type === 'identifier') {
858
+ const name = func.text;
859
+ // Only uppercase-first names (constructor heuristic)
860
+ if (name[0] >= 'A' && name[0] <= 'Z') return name;
861
+ }
862
+ return null;
863
+ }
864
+
865
+ // Boolean fallback: param or ClassName(...)
866
+ if (node.type === 'boolean_operator') {
867
+ // Check operator is 'or'
868
+ const op = node.child(1);
869
+ if (op?.text === 'or') {
870
+ const right = node.child(2);
871
+ return extractConstructorName(right);
872
+ }
873
+ }
874
+
875
+ // Conditional expression: expr if cond else ClassName(...)
876
+ if (node.type === 'conditional_expression') {
877
+ // Children: [0]=truthy, [1]='if', [2]=condition, [3]='else', [4]=else_value
878
+ // Try else branch first (usually has the constructor fallback)
879
+ const elseVal = node.child(4);
880
+ const fromElse = extractConstructorName(elseVal);
881
+ if (fromElse) return fromElse;
882
+ // Also try truthy branch
883
+ const truthyVal = node.child(0);
884
+ return extractConstructorName(truthyVal);
885
+ }
886
+
887
+ // Parenthesized expression
888
+ if (node.type === 'parenthesized_expression') {
889
+ for (let i = 0; i < node.childCount; i++) {
890
+ const child = node.child(i);
891
+ if (child.type !== '(' && child.type !== ')') {
892
+ return extractConstructorName(child);
893
+ }
894
+ }
895
+ }
896
+
897
+ return null;
898
+ }
899
+
702
900
  module.exports = {
703
901
  findFunctions,
704
902
  findClasses,
@@ -707,5 +905,6 @@ module.exports = {
707
905
  findImportsInCode,
708
906
  findExportsInCode,
709
907
  findUsagesInCode,
908
+ findInstanceAttributeTypes,
710
909
  parse
711
910
  };
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "ucn",
3
- "version": "3.4.0",
3
+ "version": "3.4.2",
4
4
  "description": "Code navigation built by AI, for AI. Reduces context usage when working with large codebases.",
5
5
  "main": "index.js",
6
6
  "bin": {
@@ -5853,5 +5853,386 @@ config = load_config()
5853
5853
  });
5854
5854
  });
5855
5855
 
5856
+ describe('Regression: Python self.attr.method() resolution', () => {
5857
+ it('findCallsInCode should detect selfAttribute for self.X.method()', () => {
5858
+ const { getParser, getLanguageModule } = require('../languages');
5859
+ const parser = getParser('python');
5860
+ const langModule = getLanguageModule('python');
5861
+
5862
+ const code = `class ReportGenerator:
5863
+ def __init__(self, analyzer):
5864
+ self.analyzer = analyzer
5865
+ self.name = "test"
5866
+
5867
+ def generate(self):
5868
+ result = self.analyzer.analyze(data)
5869
+ self.save(result)
5870
+ helper(result)
5871
+ self.name.upper()
5872
+ `;
5873
+ const calls = langModule.findCallsInCode(code, parser);
5874
+
5875
+ // self.analyzer.analyze() should have selfAttribute
5876
+ const analyzeCall = calls.find(c => c.name === 'analyze');
5877
+ assert.ok(analyzeCall, 'Should find analyze call');
5878
+ assert.strictEqual(analyzeCall.selfAttribute, 'analyzer');
5879
+ assert.strictEqual(analyzeCall.receiver, 'self');
5880
+ assert.strictEqual(analyzeCall.isMethod, true);
5881
+
5882
+ // self.save() should NOT have selfAttribute (direct self method)
5883
+ const saveCall = calls.find(c => c.name === 'save');
5884
+ assert.ok(saveCall, 'Should find save call');
5885
+ assert.strictEqual(saveCall.selfAttribute, undefined);
5886
+ assert.strictEqual(saveCall.receiver, 'self');
5887
+
5888
+ // helper() should be a regular function call
5889
+ const helperCall = calls.find(c => c.name === 'helper');
5890
+ assert.ok(helperCall, 'Should find helper call');
5891
+ assert.strictEqual(helperCall.isMethod, false);
5892
+ assert.strictEqual(helperCall.selfAttribute, undefined);
5893
+
5894
+ // self.name.upper() should have selfAttribute but string method
5895
+ const upperCall = calls.find(c => c.name === 'upper');
5896
+ assert.ok(upperCall, 'Should find upper call');
5897
+ assert.strictEqual(upperCall.selfAttribute, 'name');
5898
+ });
5899
+
5900
+ it('findInstanceAttributeTypes should parse __init__ assignments', () => {
5901
+ const { getParser, getLanguageModule } = require('../languages');
5902
+ const parser = getParser('python');
5903
+ const langModule = getLanguageModule('python');
5904
+
5905
+ const code = `class ReportGenerator:
5906
+ def __init__(self, analyzer=None, db=None):
5907
+ self.analyzer = InstrumentAnalyzer(config)
5908
+ self.db = db or DatabaseClient()
5909
+ self.name = "test"
5910
+ self.count = 0
5911
+ self.items = []
5912
+ self.scanner = (param or MarketScanner())
5913
+
5914
+ class OtherClass:
5915
+ def __init__(self):
5916
+ self.helper = HelperTool()
5917
+ `;
5918
+ const result = langModule.findInstanceAttributeTypes(code, parser);
5919
+
5920
+ // ReportGenerator
5921
+ const rg = result.get('ReportGenerator');
5922
+ assert.ok(rg, 'Should find ReportGenerator');
5923
+ assert.strictEqual(rg.get('analyzer'), 'InstrumentAnalyzer');
5924
+ assert.strictEqual(rg.get('db'), 'DatabaseClient');
5925
+ assert.strictEqual(rg.get('scanner'), 'MarketScanner');
5926
+ assert.strictEqual(rg.has('name'), false, 'Should skip string literals');
5927
+ assert.strictEqual(rg.has('count'), false, 'Should skip number literals');
5928
+ assert.strictEqual(rg.has('items'), false, 'Should skip list literals');
5929
+
5930
+ // OtherClass
5931
+ const oc = result.get('OtherClass');
5932
+ assert.ok(oc, 'Should find OtherClass');
5933
+ assert.strictEqual(oc.get('helper'), 'HelperTool');
5934
+ });
5935
+
5936
+ it('findCallees should resolve self.attr.method() to target class', () => {
5937
+ const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ucn-selfattr-'));
5938
+ try {
5939
+ fs.writeFileSync(path.join(tmpDir, 'analyzer.py'), `class InstrumentAnalyzer:
5940
+ def __init__(self, config):
5941
+ self.config = config
5942
+
5943
+ def analyze_instrument(self, data):
5944
+ return process(data)
5945
+
5946
+ def get_summary(self):
5947
+ return "summary"
5948
+ `);
5949
+ fs.writeFileSync(path.join(tmpDir, 'report.py'), `from analyzer import InstrumentAnalyzer
5950
+
5951
+ class ReportGenerator:
5952
+ def __init__(self, config):
5953
+ self.analyzer = InstrumentAnalyzer(config)
5954
+
5955
+ def generate_report(self, data):
5956
+ result = self.analyzer.analyze_instrument(data)
5957
+ summary = self.analyzer.get_summary()
5958
+ return format_output(result, summary)
5959
+
5960
+ def format_output(result, summary):
5961
+ return str(result) + summary
5962
+ `);
5963
+
5964
+ const index = new ProjectIndex(tmpDir);
5965
+ index.build('**/*.py', { quiet: true });
5966
+
5967
+ // Find generate_report definition
5968
+ const defs = index.symbols.get('generate_report');
5969
+ assert.ok(defs && defs.length > 0, 'Should find generate_report');
5970
+
5971
+ const callees = index.findCallees(defs[0]);
5972
+ const calleeNames = callees.map(c => c.name);
5973
+
5974
+ assert.ok(calleeNames.includes('analyze_instrument'),
5975
+ `Should resolve self.analyzer.analyze_instrument(), got: ${calleeNames.join(', ')}`);
5976
+ assert.ok(calleeNames.includes('get_summary'),
5977
+ `Should resolve self.analyzer.get_summary(), got: ${calleeNames.join(', ')}`);
5978
+ assert.ok(calleeNames.includes('format_output'),
5979
+ `Should include direct call format_output(), got: ${calleeNames.join(', ')}`);
5980
+
5981
+ // Verify the resolved callee points to InstrumentAnalyzer's method
5982
+ const analyzeCallee = callees.find(c => c.name === 'analyze_instrument');
5983
+ assert.strictEqual(analyzeCallee.className, 'InstrumentAnalyzer',
5984
+ 'Resolved callee should belong to InstrumentAnalyzer class');
5985
+ } finally {
5986
+ fs.rmSync(tmpDir, { recursive: true, force: true });
5987
+ }
5988
+ });
5989
+
5990
+ it('findCallers should find callers through self.attr.method()', () => {
5991
+ const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ucn-selfattr-callers-'));
5992
+ try {
5993
+ fs.writeFileSync(path.join(tmpDir, 'analyzer.py'), `class InstrumentAnalyzer:
5994
+ def __init__(self, config):
5995
+ self.config = config
5996
+
5997
+ def analyze_instrument(self, data):
5998
+ return process(data)
5999
+ `);
6000
+ fs.writeFileSync(path.join(tmpDir, 'report.py'), `from analyzer import InstrumentAnalyzer
6001
+
6002
+ class ReportGenerator:
6003
+ def __init__(self, config):
6004
+ self.analyzer = InstrumentAnalyzer(config)
6005
+
6006
+ def generate_report(self, data):
6007
+ result = self.analyzer.analyze_instrument(data)
6008
+ return result
6009
+ `);
6010
+
6011
+ const index = new ProjectIndex(tmpDir);
6012
+ index.build('**/*.py', { quiet: true });
6013
+
6014
+ const callers = index.findCallers('analyze_instrument');
6015
+ assert.ok(callers.length >= 1,
6016
+ `Should find at least 1 caller for analyze_instrument, got ${callers.length}`);
6017
+
6018
+ const reportCaller = callers.find(c => c.callerName === 'generate_report');
6019
+ assert.ok(reportCaller,
6020
+ `Should find generate_report as caller, got: ${callers.map(c => c.callerName).join(', ')}`);
6021
+ } finally {
6022
+ fs.rmSync(tmpDir, { recursive: true, force: true });
6023
+ }
6024
+ });
6025
+ });
6026
+
6027
+ describe('Regression: Python decorated function callees', () => {
6028
+ it('findCallees should work for @property and other decorated methods', () => {
6029
+ const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ucn-decorated-'));
6030
+ try {
6031
+ fs.writeFileSync(path.join(tmpDir, 'models.py'), `class Headers:
6032
+ def get(self, key):
6033
+ return self.data.get(key)
6034
+
6035
+ class Response:
6036
+ def __init__(self, headers):
6037
+ self.headers = Headers(headers)
6038
+
6039
+ @property
6040
+ def charset_encoding(self):
6041
+ content_type = self.headers.get("Content-Type")
6042
+ return parse_charset(content_type)
6043
+
6044
+ @staticmethod
6045
+ def create(data):
6046
+ return Response(data)
6047
+
6048
+ def normal_method(self):
6049
+ return self.headers.get("Accept")
6050
+ `);
6051
+
6052
+ const index = new ProjectIndex(tmpDir);
6053
+ index.build('**/*.py', { quiet: true });
6054
+
6055
+ // @property method should have callees
6056
+ const propDefs = index.symbols.get('charset_encoding');
6057
+ assert.ok(propDefs && propDefs.length > 0, 'Should find charset_encoding');
6058
+ const propCallees = index.findCallees(propDefs[0]);
6059
+ const propCalleeNames = propCallees.map(c => c.name);
6060
+ assert.ok(propCalleeNames.includes('parse_charset') || propCalleeNames.includes('get'),
6061
+ `@property should have callees, got: ${propCalleeNames.join(', ')}`);
6062
+
6063
+ // Normal method should also have callees
6064
+ const normalDefs = index.symbols.get('normal_method');
6065
+ assert.ok(normalDefs && normalDefs.length > 0, 'Should find normal_method');
6066
+ const normalCallees = index.findCallees(normalDefs[0]);
6067
+ assert.ok(normalCallees.length > 0,
6068
+ `Normal method should have callees, got: ${normalCallees.map(c => c.name).join(', ')}`);
6069
+ } finally {
6070
+ fs.rmSync(tmpDir, { recursive: true, force: true });
6071
+ }
6072
+ });
6073
+ });
6074
+
6075
+ describe('Regression: Python conditional expression in attribute types', () => {
6076
+ it('findInstanceAttributeTypes should resolve conditional expressions', () => {
6077
+ const Parser = require('tree-sitter');
6078
+ const Python = require('tree-sitter-python');
6079
+ const parser = new Parser();
6080
+ parser.setLanguage(Python);
6081
+ const pythonParser = require('../languages/python');
6082
+
6083
+ const code = `
6084
+ class Live:
6085
+ def __init__(self, renderable=None, console=None):
6086
+ self._live_render = renderable if renderable else LiveRender(Text())
6087
+ self.console = console or Console()
6088
+ self.plain = "hello"
6089
+ `;
6090
+ const result = pythonParser.findInstanceAttributeTypes(code, parser);
6091
+ assert.ok(result.has('Live'), 'Should find Live class');
6092
+ const attrs = result.get('Live');
6093
+ assert.strictEqual(attrs.get('_live_render'), 'LiveRender', 'Should resolve conditional to LiveRender');
6094
+ assert.strictEqual(attrs.get('console'), 'Console', 'Should still resolve boolean or pattern');
6095
+ assert.ok(!attrs.has('plain'), 'Should skip string literals');
6096
+ });
6097
+ });
6098
+
6099
+ describe('Regression: Python @dataclass field annotation types', () => {
6100
+ it('findInstanceAttributeTypes should extract types from @dataclass annotated fields', () => {
6101
+ const Parser = require('tree-sitter');
6102
+ const Python = require('tree-sitter-python');
6103
+ const parser = new Parser();
6104
+ parser.setLanguage(Python);
6105
+ const pythonParser = require('../languages/python');
6106
+
6107
+ const code = `
6108
+ from dataclasses import dataclass, field
6109
+
6110
+ @dataclass
6111
+ class Line:
6112
+ depth: int = 0
6113
+ bracket_tracker: BracketTracker = field(default_factory=BracketTracker)
6114
+ inside_brackets: bool = False
6115
+ comments: list = field(default_factory=list)
6116
+ mode: Mode = Mode.DEFAULT
6117
+ `;
6118
+ const result = pythonParser.findInstanceAttributeTypes(code, parser);
6119
+ assert.ok(result.has('Line'), 'Should find Line class');
6120
+ const attrs = result.get('Line');
6121
+ assert.strictEqual(attrs.get('bracket_tracker'), 'BracketTracker', 'Should extract BracketTracker from annotation');
6122
+ assert.strictEqual(attrs.get('mode'), 'Mode', 'Should extract Mode from annotation');
6123
+ assert.ok(!attrs.has('depth'), 'Should skip int primitive');
6124
+ assert.ok(!attrs.has('inside_brackets'), 'Should skip bool primitive');
6125
+ assert.ok(!attrs.has('comments'), 'Should skip list primitive');
6126
+ });
6127
+
6128
+ it('findInstanceAttributeTypes should not scan non-dataclass classes for field annotations', () => {
6129
+ const Parser = require('tree-sitter');
6130
+ const Python = require('tree-sitter-python');
6131
+ const parser = new Parser();
6132
+ parser.setLanguage(Python);
6133
+ const pythonParser = require('../languages/python');
6134
+
6135
+ const code = `
6136
+ class RegularClass:
6137
+ name: str = "default"
6138
+ tracker: BracketTracker = None
6139
+
6140
+ def __init__(self):
6141
+ self.helper = Helper()
6142
+ `;
6143
+ const result = pythonParser.findInstanceAttributeTypes(code, parser);
6144
+ // Should find Helper from __init__ but NOT BracketTracker from class-level annotation
6145
+ assert.ok(result.has('RegularClass'), 'Should find RegularClass');
6146
+ const attrs = result.get('RegularClass');
6147
+ assert.strictEqual(attrs.get('helper'), 'Helper', 'Should find Helper from __init__');
6148
+ assert.ok(!attrs.has('tracker'), 'Should NOT extract from non-dataclass class annotations');
6149
+ });
6150
+
6151
+ it('findCallees should resolve @dataclass attribute method calls', () => {
6152
+ const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ucn-dataclass-'));
6153
+ try {
6154
+ fs.writeFileSync(path.join(tmpDir, 'linegen.py'), `
6155
+ from dataclasses import dataclass, field
6156
+
6157
+ class BracketTracker:
6158
+ def any_open_brackets(self):
6159
+ return len(self.brackets) > 0
6160
+
6161
+ def mark(self, leaf):
6162
+ self.brackets.append(leaf)
6163
+
6164
+ @dataclass
6165
+ class Line:
6166
+ bracket_tracker: BracketTracker = field(default_factory=BracketTracker)
6167
+
6168
+ def should_split(self):
6169
+ if self.bracket_tracker.any_open_brackets():
6170
+ return True
6171
+ self.bracket_tracker.mark(None)
6172
+ return False
6173
+ `);
6174
+ const index = new ProjectIndex(tmpDir);
6175
+ index.build('**/*.py', { quiet: true });
6176
+
6177
+ const defs = index.symbols.get('should_split');
6178
+ assert.ok(defs && defs.length > 0, 'Should find should_split');
6179
+ const callees = index.findCallees(defs[0]);
6180
+ const calleeNames = callees.map(c => c.name);
6181
+ assert.ok(calleeNames.includes('any_open_brackets'),
6182
+ `Should resolve self.bracket_tracker.any_open_brackets(), got: ${calleeNames.join(', ')}`);
6183
+ assert.ok(calleeNames.includes('mark'),
6184
+ `Should resolve self.bracket_tracker.mark(), got: ${calleeNames.join(', ')}`);
6185
+ } finally {
6186
+ fs.rmSync(tmpDir, { recursive: true, force: true });
6187
+ }
6188
+ });
6189
+ });
6190
+
6191
+ describe('Regression: Python self.method() same-class resolution', () => {
6192
+ it('findCallees should resolve self.method() to same-class methods', () => {
6193
+ const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ucn-selfmethod-'));
6194
+ try {
6195
+ fs.writeFileSync(path.join(tmpDir, 'service.py'), `
6196
+ class DataService:
6197
+ def _fetch_remote(self, key, days):
6198
+ return self._make_request(f"/api/{key}")
6199
+
6200
+ def _make_request(self, url):
6201
+ return None
6202
+
6203
+ def get_records(self, key, days=365):
6204
+ if self._is_valid(key):
6205
+ return self._fetch_remote(key, days)
6206
+ return None
6207
+
6208
+ def _is_valid(self, key):
6209
+ return len(key) > 0
6210
+ `);
6211
+ const index = new ProjectIndex(tmpDir);
6212
+ index.build('**/*.py', { quiet: true });
6213
+
6214
+ // get_records should have _fetch_remote and _is_valid as callees
6215
+ const defs = index.symbols.get('get_records');
6216
+ assert.ok(defs && defs.length > 0, 'Should find get_records');
6217
+ const callees = index.findCallees(defs[0]);
6218
+ const calleeNames = callees.map(c => c.name);
6219
+ assert.ok(calleeNames.includes('_fetch_remote'),
6220
+ `Should resolve self._fetch_remote(), got: ${calleeNames.join(', ')}`);
6221
+ assert.ok(calleeNames.includes('_is_valid'),
6222
+ `Should resolve self._is_valid(), got: ${calleeNames.join(', ')}`);
6223
+
6224
+ // _fetch_remote should have get_records as caller
6225
+ const fetchDefs = index.symbols.get('_fetch_remote');
6226
+ assert.ok(fetchDefs && fetchDefs.length > 0, 'Should find _fetch_remote');
6227
+ const callers = index.findCallers('_fetch_remote');
6228
+ const callerNames = callers.map(c => c.callerName);
6229
+ assert.ok(callerNames.includes('get_records'),
6230
+ `Should find get_records as caller of _fetch_remote, got: ${callerNames.join(', ')}`);
6231
+ } finally {
6232
+ fs.rmSync(tmpDir, { recursive: true, force: true });
6233
+ }
6234
+ });
6235
+ });
6236
+
5856
6237
  console.log('UCN v3 Test Suite');
5857
6238
  console.log('Run with: node --test test/parser.test.js');