ucn 3.4.0 → 3.4.1

package/core/project.js

@@ -1009,11 +1009,25 @@ class ProjectIndex {
 
                     // Smart method call handling
                     if (call.isMethod) {
-                        // Always skip this/self/cls calls (internal state access, not function calls)
-                        if (['this', 'self', 'cls'].includes(call.receiver)) continue;
-                        // Go doesn't use this/self/cls - always include Go method calls
-                        // For other languages, skip method calls unless explicitly requested
-                        if (fileEntry.language !== 'go' && !options.includeMethods) continue;
+                        if (call.selfAttribute && fileEntry.language === 'python') {
+                            // self.attr.method() — resolve via attribute type inference
+                            const callerSymbol = this.findEnclosingFunction(filePath, call.line, true);
+                            if (!callerSymbol?.className) continue;
+                            const attrTypes = this.getInstanceAttributeTypes(filePath, callerSymbol.className);
+                            if (!attrTypes) continue;
+                            const targetClass = attrTypes.get(call.selfAttribute);
+                            if (!targetClass) continue;
+                            // Check if any definition of searched function belongs to targetClass
+                            const matchesDef = definitions.some(d => d.className === targetClass);
+                            if (!matchesDef) continue;
+                            // Falls through to add as caller
+                        } else {
+                            // Always skip this/self/cls calls (internal state access, not function calls)
+                            if (['this', 'self', 'cls'].includes(call.receiver)) continue;
+                            // Go doesn't use this/self/cls - always include Go method calls
+                            // For other languages, skip method calls unless explicitly requested
+                            if (fileEntry.language !== 'go' && !options.includeMethods) continue;
+                        }
                     }
 
                     // Skip definition lines
@@ -1110,6 +1124,7 @@ class ProjectIndex {
             const language = fileEntry?.language;
 
             const callees = new Map();  // key -> { name, bindingId, count }
+            let selfAttrCalls = null;   // collected for Python self.attr.method() resolution
 
             for (const call of calls) {
                 // Filter to calls within this function's scope using enclosingFunction
@@ -1119,14 +1134,26 @@ class ProjectIndex {
 
                 // Smart method call handling:
                 // - Go: include all method calls (Go doesn't use this/self/cls)
+                // - Python self.attr.method(): resolve via selfAttribute (handled below)
                 // - Other languages: skip method calls unless explicitly requested
                 if (call.isMethod) {
-                    if (language !== 'go' && !options.includeMethods) continue;
+                    if (call.selfAttribute && language === 'python') {
+                        // Will be resolved in second pass below
+                    } else if (language !== 'go' && !options.includeMethods) {
+                        continue;
+                    }
                 }
 
                 // Skip keywords and built-ins
                 if (this.isKeyword(call.name, language)) continue;
 
+                // Collect selfAttribute calls for second-pass resolution
+                if (call.selfAttribute && language === 'python') {
+                    if (!selfAttrCalls) selfAttrCalls = [];
+                    selfAttrCalls.push(call);
+                    continue;
+                }
+
                 // Resolve binding within this file (without mutating cached call objects)
                 let calleeKey = call.bindingId || call.name;
                 let bindingResolved = call.bindingId;
@@ -1179,6 +1206,36 @@ class ProjectIndex {
                 }
             }
 
+            // Second pass: resolve Python self.attr.method() calls
+            if (selfAttrCalls && def.className) {
+                const attrTypes = this.getInstanceAttributeTypes(def.file, def.className);
+                if (attrTypes) {
+                    for (const call of selfAttrCalls) {
+                        const targetClass = attrTypes.get(call.selfAttribute);
+                        if (!targetClass) continue;
+
+                        // Find method in symbol table where className matches
+                        const symbols = this.symbols.get(call.name);
+                        if (!symbols) continue;
+
+                        const match = symbols.find(s => s.className === targetClass);
+                        if (!match) continue;
+
+                        const key = match.bindingId || `${targetClass}.${call.name}`;
+                        const existing = callees.get(key);
+                        if (existing) {
+                            existing.count += 1;
+                        } else {
+                            callees.set(key, {
+                                name: call.name,
+                                bindingId: match.bindingId,
+                                count: 1
+                            });
+                        }
+                    }
+                }
+            }
+
             // Look up each callee in the symbol table
             // For methods, prefer callees from: 1) same file, 2) same package, 3) same receiver type
             const result = [];
@@ -1490,8 +1547,9 @@ class ProjectIndex {
         const fileEntry = this.files.get(filePath);
         if (!fileEntry) return null;
 
+        const nonCallableTypes = new Set(['class', 'struct', 'interface', 'type', 'state']);
         for (const symbol of fileEntry.symbols) {
-            if (symbol.type === 'function' &&
+            if (!nonCallableTypes.has(symbol.type) &&
                 symbol.startLine <= lineNum &&
                 symbol.endLine >= lineNum) {
                 if (returnSymbol) {
@@ -1503,6 +1561,35 @@ class ProjectIndex {
         return null;
     }
 
+    /**
+     * Get instance attribute types for a class in a file.
+     * Returns Map<attrName, typeName> for a given className.
+     * Caches results per file.
+     */
+    getInstanceAttributeTypes(filePath, className) {
+        if (!this._attrTypeCache) this._attrTypeCache = new Map();
+
+        let fileCache = this._attrTypeCache.get(filePath);
+        if (!fileCache) {
+            const fileEntry = this.files.get(filePath);
+            if (!fileEntry || fileEntry.language !== 'python') return null;
+
+            const langModule = getLanguageModule('python');
+            if (!langModule?.findInstanceAttributeTypes) return null;
+
+            try {
+                const content = fs.readFileSync(filePath, 'utf-8');
+                const parser = getParser('python');
+                fileCache = langModule.findInstanceAttributeTypes(content, parser);
+                this._attrTypeCache.set(filePath, fileCache);
+            } catch {
+                return null;
+            }
+        }
+
+        return fileCache.get(className) || null;
+    }
+
     /**
      * Extract type names from a function definition
      */

package/languages/python.js

@@ -400,9 +400,14 @@ function findCallsInCode(code, parser) {
     traverseTree(tree.rootNode, (node) => {
         // Track function entry
         if (isFunctionNode(node)) {
+            // Use decorated_definition start line if present, to match symbol index
+            let startLine = node.startPosition.row + 1;
+            if (node.parent && node.parent.type === 'decorated_definition') {
+                startLine = node.parent.startPosition.row + 1;
+            }
             functionStack.push({
                 name: extractFunctionName(node),
-                startLine: node.startPosition.row + 1,
+                startLine,
                 endLine: node.endPosition.row + 1
             });
         }
@@ -425,16 +430,32 @@ function findCallsInCode(code, parser) {
                     uncertain
                 });
             } else if (funcNode.type === 'attribute') {
-                // Method/attribute call: obj.foo()
+                // Method/attribute call: obj.foo() or self.attr.foo()
                 const attrNode = funcNode.childForFieldName('attribute');
                 const objNode = funcNode.childForFieldName('object');
 
                 if (attrNode) {
+                    let receiver = objNode?.type === 'identifier' ? objNode.text : undefined;
+                    let selfAttribute = undefined;
+
+                    // Detect self.X.method() pattern: objNode is attribute access on self/cls
+                    if (objNode?.type === 'attribute') {
+                        const innerObj = objNode.childForFieldName('object');
+                        const innerAttr = objNode.childForFieldName('attribute');
+                        if (innerObj?.type === 'identifier' &&
+                            ['self', 'cls'].includes(innerObj.text) &&
+                            innerAttr) {
+                            selfAttribute = innerAttr.text;
+                            receiver = innerObj.text;
+                        }
+                    }
+
                     calls.push({
                         name: attrNode.text,
                         line: node.startPosition.row + 1,
                         isMethod: true,
-                        receiver: objNode?.type === 'identifier' ? objNode.text : undefined,
+                        receiver,
+                        ...(selfAttribute && { selfAttribute }),
                         enclosingFunction,
                         uncertain
                     });
@@ -699,6 +720,183 @@ function findUsagesInCode(code, name, parser) {
     return usages;
 }
 
+/**
+ * Find instance attribute types from __init__ constructor assignments.
+ * Parses self.X = ClassName(...) patterns in __init__ methods.
+ * @param {string} code - Source code to analyze
+ * @param {object} parser - Tree-sitter parser instance
+ * @returns {Map<string, Map<string, string>>} className -> (attrName -> typeName)
+ */
+function findInstanceAttributeTypes(code, parser) {
+    const tree = parseTree(parser, code);
+    const result = new Map(); // className -> Map(attrName -> typeName)
+
+    const PRIMITIVE_TYPES = new Set(['int', 'float', 'str', 'bool', 'bytes', 'list', 'dict', 'set', 'tuple', 'None', 'Any', 'object']);
+
+    traverseTree(tree.rootNode, (node) => {
+        if (node.type !== 'class_definition') return true;
+
+        const classNameNode = node.childForFieldName('name');
+        if (!classNameNode) return true;
+        const className = classNameNode.text;
+
+        const body = node.childForFieldName('body');
+        if (!body) return false;
+
+        const attrTypes = new Map();
+
+        // Check for @dataclass decorator — scan annotated class-level fields
+        const parentNode = node.parent;
+        if (parentNode?.type === 'decorated_definition') {
+            for (let d = 0; d < parentNode.childCount; d++) {
+                const dec = parentNode.child(d);
+                if (dec.type !== 'decorator') continue;
+                // Match @dataclass or @dataclasses.dataclass
+                const decText = dec.text;
+                if (decText.startsWith('@dataclass') || decText.includes('.dataclass')) {
+                    // Scan class body for annotated fields: name: Type = ...
+                    for (let i = 0; i < body.childCount; i++) {
+                        const stmt = body.child(i);
+                        if (stmt.type !== 'expression_statement') continue;
+                        const assign = stmt.firstChild;
+                        if (!assign || assign.type !== 'assignment') continue;
+
+                        // Must have a type annotation
+                        const typeNode = assign.childForFieldName('type');
+                        if (!typeNode) continue;
+
+                        // Extract type name from annotation
+                        const typeIdent = typeNode.type === 'type' ? typeNode.firstChild : typeNode;
+                        if (!typeIdent || typeIdent.type !== 'identifier') continue;
+                        const typeName = typeIdent.text;
+
+                        // Skip primitives and lowercase types
+                        if (PRIMITIVE_TYPES.has(typeName)) continue;
+                        if (typeName[0] < 'A' || typeName[0] > 'Z') continue;
+
+                        // Field name from LHS
+                        const lhs = assign.childForFieldName('left');
+                        if (!lhs || lhs.type !== 'identifier') continue;
+                        attrTypes.set(lhs.text, typeName);
+                    }
+                    break;
+                }
+            }
+        }
+
+        // Scan __init__ for self.X = ClassName(...) assignments
+        for (let i = 0; i < body.childCount; i++) {
+            let child = body.child(i);
+            // Handle decorated_definition wrapper
+            if (child.type === 'decorated_definition') {
+                for (let j = 0; j < child.childCount; j++) {
+                    if (child.child(j).type === 'function_definition') {
+                        child = child.child(j);
+                        break;
+                    }
+                }
+            }
+            if (child.type !== 'function_definition') continue;
+
+            const fnName = child.childForFieldName('name');
+            if (!fnName || fnName.text !== '__init__') continue;
+
+            // Found __init__, now scan for self.X = ClassName(...) assignments
+            const initBody = child.childForFieldName('body');
+            if (!initBody) continue;
+
+            traverseTree(initBody, (stmt) => {
+                if (stmt.type !== 'expression_statement') return true;
+
+                const assign = stmt.firstChild;
+                if (!assign || assign.type !== 'assignment') return true;
+
+                // LHS: self.X
+                const lhs = assign.childForFieldName('left');
+                if (!lhs || lhs.type !== 'attribute') return true;
+                const lhsObj = lhs.childForFieldName('object');
+                const lhsAttr = lhs.childForFieldName('attribute');
+                if (!lhsObj || lhsObj.text !== 'self' || !lhsAttr) return true;
+
+                const attrName = lhsAttr.text;
+
+                // RHS: ClassName(...) or param or ClassName(...)
+                const rhs = assign.childForFieldName('right');
+                if (!rhs) return true;
+
+                const typeName = extractConstructorName(rhs);
+                if (typeName) {
+                    attrTypes.set(attrName, typeName);
+                }
+
+                return true;
+            });
+        }
+
+        if (attrTypes.size > 0) {
+            result.set(className, attrTypes);
+        }
+
+        return false; // don't descend into nested classes from traverseTree
+    });
+
+    return result;
+}
+
+/**
+ * Extract constructor class name from an expression node.
+ * Handles: ClassName(...), param or ClassName(...), (param or ClassName(...)),
+ *          expr if cond else ClassName(...)
+ */
+function extractConstructorName(node) {
+    if (!node) return null;
+
+    // Direct call: ClassName(...)
+    if (node.type === 'call') {
+        const func = node.childForFieldName('function');
+        if (func?.type === 'identifier') {
+            const name = func.text;
+            // Only uppercase-first names (constructor heuristic)
+            if (name[0] >= 'A' && name[0] <= 'Z') return name;
+        }
+        return null;
+    }
+
+    // Boolean fallback: param or ClassName(...)
+    if (node.type === 'boolean_operator') {
+        // Check operator is 'or'
+        const op = node.child(1);
+        if (op?.text === 'or') {
+            const right = node.child(2);
+            return extractConstructorName(right);
+        }
+    }
+
+    // Conditional expression: expr if cond else ClassName(...)
+    if (node.type === 'conditional_expression') {
+        // Children: [0]=truthy, [1]='if', [2]=condition, [3]='else', [4]=else_value
+        // Try else branch first (usually has the constructor fallback)
+        const elseVal = node.child(4);
+        const fromElse = extractConstructorName(elseVal);
+        if (fromElse) return fromElse;
+        // Also try truthy branch
+        const truthyVal = node.child(0);
+        return extractConstructorName(truthyVal);
+    }
+
+    // Parenthesized expression
+    if (node.type === 'parenthesized_expression') {
+        for (let i = 0; i < node.childCount; i++) {
+            const child = node.child(i);
+            if (child.type !== '(' && child.type !== ')') {
+                return extractConstructorName(child);
+            }
+        }
+    }
+
+    return null;
+}
+
 module.exports = {
     findFunctions,
     findClasses,
@@ -707,5 +905,6 @@ module.exports = {
     findImportsInCode,
     findExportsInCode,
     findUsagesInCode,
+    findInstanceAttributeTypes,
     parse
 };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "ucn",
-  "version": "3.4.0",
+  "version": "3.4.1",
   "description": "Code navigation built by AI, for AI. Reduces context usage when working with large codebases.",
   "main": "index.js",
   "bin": {

package/test/parser.test.js

@@ -5853,5 +5853,340 @@ config = load_config()
     });
 });
 
+describe('Regression: Python self.attr.method() resolution', () => {
+    it('findCallsInCode should detect selfAttribute for self.X.method()', () => {
+        const { getParser, getLanguageModule } = require('../languages');
+        const parser = getParser('python');
+        const langModule = getLanguageModule('python');
+
+        const code = `class ReportGenerator:
+    def __init__(self, analyzer):
+        self.analyzer = analyzer
+        self.name = "test"
+
+    def generate(self):
+        result = self.analyzer.analyze(data)
+        self.save(result)
+        helper(result)
+        self.name.upper()
+`;
+        const calls = langModule.findCallsInCode(code, parser);
+
+        // self.analyzer.analyze() should have selfAttribute
+        const analyzeCall = calls.find(c => c.name === 'analyze');
+        assert.ok(analyzeCall, 'Should find analyze call');
+        assert.strictEqual(analyzeCall.selfAttribute, 'analyzer');
+        assert.strictEqual(analyzeCall.receiver, 'self');
+        assert.strictEqual(analyzeCall.isMethod, true);
+
+        // self.save() should NOT have selfAttribute (direct self method)
+        const saveCall = calls.find(c => c.name === 'save');
+        assert.ok(saveCall, 'Should find save call');
+        assert.strictEqual(saveCall.selfAttribute, undefined);
+        assert.strictEqual(saveCall.receiver, 'self');
+
+        // helper() should be a regular function call
+        const helperCall = calls.find(c => c.name === 'helper');
+        assert.ok(helperCall, 'Should find helper call');
+        assert.strictEqual(helperCall.isMethod, false);
+        assert.strictEqual(helperCall.selfAttribute, undefined);
+
+        // self.name.upper() should have selfAttribute but string method
+        const upperCall = calls.find(c => c.name === 'upper');
+        assert.ok(upperCall, 'Should find upper call');
+        assert.strictEqual(upperCall.selfAttribute, 'name');
+    });
+
+    it('findInstanceAttributeTypes should parse __init__ assignments', () => {
+        const { getParser, getLanguageModule } = require('../languages');
+        const parser = getParser('python');
+        const langModule = getLanguageModule('python');
+
+        const code = `class ReportGenerator:
+    def __init__(self, analyzer=None, db=None):
+        self.analyzer = InstrumentAnalyzer(config)
+        self.db = db or DatabaseClient()
+        self.name = "test"
+        self.count = 0
+        self.items = []
+        self.scanner = (param or MarketScanner())
+
+class OtherClass:
+    def __init__(self):
+        self.helper = HelperTool()
+`;
+        const result = langModule.findInstanceAttributeTypes(code, parser);
+
+        // ReportGenerator
+        const rg = result.get('ReportGenerator');
+        assert.ok(rg, 'Should find ReportGenerator');
+        assert.strictEqual(rg.get('analyzer'), 'InstrumentAnalyzer');
+        assert.strictEqual(rg.get('db'), 'DatabaseClient');
+        assert.strictEqual(rg.get('scanner'), 'MarketScanner');
+        assert.strictEqual(rg.has('name'), false, 'Should skip string literals');
+        assert.strictEqual(rg.has('count'), false, 'Should skip number literals');
+        assert.strictEqual(rg.has('items'), false, 'Should skip list literals');
+
+        // OtherClass
+        const oc = result.get('OtherClass');
+        assert.ok(oc, 'Should find OtherClass');
+        assert.strictEqual(oc.get('helper'), 'HelperTool');
+    });
+
+    it('findCallees should resolve self.attr.method() to target class', () => {
+        const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ucn-selfattr-'));
+        try {
+            fs.writeFileSync(path.join(tmpDir, 'analyzer.py'), `class InstrumentAnalyzer:
+    def __init__(self, config):
+        self.config = config
+
+    def analyze_instrument(self, data):
+        return process(data)
+
+    def get_summary(self):
+        return "summary"
+`);
+            fs.writeFileSync(path.join(tmpDir, 'report.py'), `from analyzer import InstrumentAnalyzer
+
+class ReportGenerator:
+    def __init__(self, config):
+        self.analyzer = InstrumentAnalyzer(config)
+
+    def generate_report(self, data):
+        result = self.analyzer.analyze_instrument(data)
+        summary = self.analyzer.get_summary()
+        return format_output(result, summary)
+
+def format_output(result, summary):
+    return str(result) + summary
+`);
+
+            const index = new ProjectIndex(tmpDir);
+            index.build('**/*.py', { quiet: true });
+
+            // Find generate_report definition
+            const defs = index.symbols.get('generate_report');
+            assert.ok(defs && defs.length > 0, 'Should find generate_report');
+
+            const callees = index.findCallees(defs[0]);
+            const calleeNames = callees.map(c => c.name);
+
+            assert.ok(calleeNames.includes('analyze_instrument'),
+                `Should resolve self.analyzer.analyze_instrument(), got: ${calleeNames.join(', ')}`);
+            assert.ok(calleeNames.includes('get_summary'),
+                `Should resolve self.analyzer.get_summary(), got: ${calleeNames.join(', ')}`);
+            assert.ok(calleeNames.includes('format_output'),
+                `Should include direct call format_output(), got: ${calleeNames.join(', ')}`);
+
+            // Verify the resolved callee points to InstrumentAnalyzer's method
+            const analyzeCallee = callees.find(c => c.name === 'analyze_instrument');
+            assert.strictEqual(analyzeCallee.className, 'InstrumentAnalyzer',
+                'Resolved callee should belong to InstrumentAnalyzer class');
+        } finally {
+            fs.rmSync(tmpDir, { recursive: true, force: true });
+        }
+    });
+
+    it('findCallers should find callers through self.attr.method()', () => {
+        const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ucn-selfattr-callers-'));
+        try {
+            fs.writeFileSync(path.join(tmpDir, 'analyzer.py'), `class InstrumentAnalyzer:
+    def __init__(self, config):
+        self.config = config
+
+    def analyze_instrument(self, data):
+        return process(data)
+`);
+            fs.writeFileSync(path.join(tmpDir, 'report.py'), `from analyzer import InstrumentAnalyzer
+
+class ReportGenerator:
+    def __init__(self, config):
+        self.analyzer = InstrumentAnalyzer(config)
+
+    def generate_report(self, data):
+        result = self.analyzer.analyze_instrument(data)
+        return result
+`);
+
+            const index = new ProjectIndex(tmpDir);
+            index.build('**/*.py', { quiet: true });
+
+            const callers = index.findCallers('analyze_instrument');
+            assert.ok(callers.length >= 1,
+                `Should find at least 1 caller for analyze_instrument, got ${callers.length}`);
+
+            const reportCaller = callers.find(c => c.callerName === 'generate_report');
+            assert.ok(reportCaller,
+                `Should find generate_report as caller, got: ${callers.map(c => c.callerName).join(', ')}`);
+        } finally {
+            fs.rmSync(tmpDir, { recursive: true, force: true });
+        }
+    });
+});
+
+describe('Regression: Python decorated function callees', () => {
+    it('findCallees should work for @property and other decorated methods', () => {
+        const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ucn-decorated-'));
+        try {
+            fs.writeFileSync(path.join(tmpDir, 'models.py'), `class Headers:
+    def get(self, key):
+        return self.data.get(key)
+
+class Response:
+    def __init__(self, headers):
+        self.headers = Headers(headers)
+
+    @property
+    def charset_encoding(self):
+        content_type = self.headers.get("Content-Type")
+        return parse_charset(content_type)
+
+    @staticmethod
+    def create(data):
+        return Response(data)
+
+    def normal_method(self):
+        return self.headers.get("Accept")
+`);
+
+            const index = new ProjectIndex(tmpDir);
+            index.build('**/*.py', { quiet: true });
+
+            // @property method should have callees
+            const propDefs = index.symbols.get('charset_encoding');
+            assert.ok(propDefs && propDefs.length > 0, 'Should find charset_encoding');
+            const propCallees = index.findCallees(propDefs[0]);
+            const propCalleeNames = propCallees.map(c => c.name);
+            assert.ok(propCalleeNames.includes('parse_charset') || propCalleeNames.includes('get'),
+                `@property should have callees, got: ${propCalleeNames.join(', ')}`);
+
+            // Normal method should also have callees
+            const normalDefs = index.symbols.get('normal_method');
+            assert.ok(normalDefs && normalDefs.length > 0, 'Should find normal_method');
+            const normalCallees = index.findCallees(normalDefs[0]);
+            assert.ok(normalCallees.length > 0,
+                `Normal method should have callees, got: ${normalCallees.map(c => c.name).join(', ')}`);
+        } finally {
+            fs.rmSync(tmpDir, { recursive: true, force: true });
+        }
+    });
+});
+
+describe('Regression: Python conditional expression in attribute types', () => {
+    it('findInstanceAttributeTypes should resolve conditional expressions', () => {
+        const Parser = require('tree-sitter');
+        const Python = require('tree-sitter-python');
+        const parser = new Parser();
+        parser.setLanguage(Python);
+        const pythonParser = require('../languages/python');
+
+        const code = `
+class Live:
+    def __init__(self, renderable=None, console=None):
+        self._live_render = renderable if renderable else LiveRender(Text())
+        self.console = console or Console()
+        self.plain = "hello"
+`;
+        const result = pythonParser.findInstanceAttributeTypes(code, parser);
+        assert.ok(result.has('Live'), 'Should find Live class');
+        const attrs = result.get('Live');
+        assert.strictEqual(attrs.get('_live_render'), 'LiveRender', 'Should resolve conditional to LiveRender');
+        assert.strictEqual(attrs.get('console'), 'Console', 'Should still resolve boolean or pattern');
+        assert.ok(!attrs.has('plain'), 'Should skip string literals');
+    });
+});
+
+describe('Regression: Python @dataclass field annotation types', () => {
+    it('findInstanceAttributeTypes should extract types from @dataclass annotated fields', () => {
+        const Parser = require('tree-sitter');
+        const Python = require('tree-sitter-python');
+        const parser = new Parser();
+        parser.setLanguage(Python);
+        const pythonParser = require('../languages/python');
+
+        const code = `
+from dataclasses import dataclass, field
+
+@dataclass
+class Line:
+    depth: int = 0
+    bracket_tracker: BracketTracker = field(default_factory=BracketTracker)
+    inside_brackets: bool = False
+    comments: list = field(default_factory=list)
+    mode: Mode = Mode.DEFAULT
+`;
+        const result = pythonParser.findInstanceAttributeTypes(code, parser);
+        assert.ok(result.has('Line'), 'Should find Line class');
+        const attrs = result.get('Line');
+        assert.strictEqual(attrs.get('bracket_tracker'), 'BracketTracker', 'Should extract BracketTracker from annotation');
+        assert.strictEqual(attrs.get('mode'), 'Mode', 'Should extract Mode from annotation');
+        assert.ok(!attrs.has('depth'), 'Should skip int primitive');
+        assert.ok(!attrs.has('inside_brackets'), 'Should skip bool primitive');
+        assert.ok(!attrs.has('comments'), 'Should skip list primitive');
+    });
+
+    it('findInstanceAttributeTypes should not scan non-dataclass classes for field annotations', () => {
+        const Parser = require('tree-sitter');
+        const Python = require('tree-sitter-python');
+        const parser = new Parser();
+        parser.setLanguage(Python);
+        const pythonParser = require('../languages/python');
+
+        const code = `
+class RegularClass:
+    name: str = "default"
+    tracker: BracketTracker = None
+
+    def __init__(self):
+        self.helper = Helper()
+`;
+        const result = pythonParser.findInstanceAttributeTypes(code, parser);
+        // Should find Helper from __init__ but NOT BracketTracker from class-level annotation
+        assert.ok(result.has('RegularClass'), 'Should find RegularClass');
+        const attrs = result.get('RegularClass');
+        assert.strictEqual(attrs.get('helper'), 'Helper', 'Should find Helper from __init__');
+        assert.ok(!attrs.has('tracker'), 'Should NOT extract from non-dataclass class annotations');
+    });
+
+    it('findCallees should resolve @dataclass attribute method calls', () => {
+        const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ucn-dataclass-'));
+        try {
+            fs.writeFileSync(path.join(tmpDir, 'linegen.py'), `
+from dataclasses import dataclass, field
+
+class BracketTracker:
+    def any_open_brackets(self):
+        return len(self.brackets) > 0
+
+    def mark(self, leaf):
+        self.brackets.append(leaf)
+
+@dataclass
+class Line:
+    bracket_tracker: BracketTracker = field(default_factory=BracketTracker)
+
+    def should_split(self):
+        if self.bracket_tracker.any_open_brackets():
+            return True
+        self.bracket_tracker.mark(None)
+        return False
+`);
+            const index = new ProjectIndex(tmpDir);
+            index.build('**/*.py', { quiet: true });
+
+            const defs = index.symbols.get('should_split');
+            assert.ok(defs && defs.length > 0, 'Should find should_split');
+            const callees = index.findCallees(defs[0]);
+            const calleeNames = callees.map(c => c.name);
+            assert.ok(calleeNames.includes('any_open_brackets'),
+                `Should resolve self.bracket_tracker.any_open_brackets(), got: ${calleeNames.join(', ')}`);
+            assert.ok(calleeNames.includes('mark'),
+                `Should resolve self.bracket_tracker.mark(), got: ${calleeNames.join(', ')}`);
+        } finally {
+            fs.rmSync(tmpDir, { recursive: true, force: true });
+        }
+    });
+});
+
 console.log('UCN v3 Test Suite');
 console.log('Run with: node --test test/parser.test.js');