repomix/lib/core/treeSitter/parseStrategies/PythonParseStrategy.js

A tool to pack repository contents to single file for AI consumption

import { BaseParseStrategy } from './BaseParseStrategy.js';
var CaptureType;
(function (CaptureType) {
    CaptureType["Comment"] = "comment";
    CaptureType["Class"] = "definition.class";
    CaptureType["Function"] = "definition.function";
    CaptureType["Docstring"] = "docstring";
    CaptureType["TypeAlias"] = "definition.type_alias";
})(CaptureType || (CaptureType = {}));
export class PythonParseStrategy extends BaseParseStrategy {
    parseCapture(capture, lines, processedChunks, _context) {
        const { node, name } = capture;
        const startRow = node.startPosition.row;
        const endRow = node.endPosition.row;
        if (!this.validateLineExists(lines, startRow)) {
            return null;
        }
        const captureTypes = this.getCaptureTypes(name, CaptureType);
        if (captureTypes.has(CaptureType.Class)) {
            return this.parseClassDefinition(lines, startRow, processedChunks).content;
        }
        if (captureTypes.has(CaptureType.Function)) {
            return this.parseFunctionDefinition(lines, startRow, processedChunks).content;
        }
        if (captureTypes.has(CaptureType.Docstring)) {
            return this.parseDocstringOrComment(lines, startRow, endRow, processedChunks).content;
        }
        if (captureTypes.has(CaptureType.Comment)) {
            return this.parseDocstringOrComment(lines, startRow, endRow, processedChunks).content;
        }
        if (captureTypes.has(CaptureType.TypeAlias)) {
            return this.parseTypeAlias(lines, startRow, processedChunks).content;
        }
        return null;
    }
    getDecorators(lines, startRow) {
        const decorators = [];
        let currentRow = startRow - 1;
        while (currentRow >= 0) {
            const line = lines[currentRow].trim();
            if (line.startsWith('@')) {
                decorators.unshift(line);
            }
            else {
                break;
            }
            currentRow--;
        }
        return decorators;
    }
    getClassInheritance(lines, startRow) {
        const line = lines[startRow];
        const match = line.match(/class\s+\w+\s*\((.*?)\):/);
        return match ? line.replace(/:\s*$/, '') : line.replace(/:\s*$/, '');
    }
    getFunctionSignature(lines, startRow) {
        const line = lines[startRow];
        const match = line.match(/def\s+(\w+)\s*\((.*?)\)(\s*->\s*[^:]+)?:/);
        if (!match)
            return null;
        return line.replace(/:\s*$/, '');
    }
    parseClassDefinition(lines, startRow, processedChunks) {
        const decorators = this.getDecorators(lines, startRow);
        const classDefinition = this.getClassInheritance(lines, startRow);
        const fullDefinition = [...decorators, classDefinition].join('\n');
        if (!this.checkAndAddToProcessed(fullDefinition, processedChunks)) {
            return this.createNullResult();
        }
        return this.createResult(fullDefinition);
    }
    parseFunctionDefinition(lines, startRow, processedChunks) {
        const decorators = this.getDecorators(lines, startRow);
        const signature = this.getFunctionSignature(lines, startRow);
        if (!signature) {
            return this.createNullResult();
        }
        const fullDefinition = [...decorators, signature].join('\n');
        if (!this.checkAndAddToProcessed(fullDefinition, processedChunks)) {
            return this.createNullResult();
        }
        return this.createResult(fullDefinition);
    }
    parseDocstringOrComment(lines, startRow, endRow, processedChunks) {
        const content = lines.slice(startRow, endRow + 1).join('\n');
        if (!this.checkAndAddToProcessed(content, processedChunks)) {
            return this.createNullResult();
        }
        return this.createResult(content);
    }
    parseTypeAlias(lines, startRow, processedChunks) {
        const typeAlias = lines[startRow].trim();
        if (!this.checkAndAddToProcessed(typeAlias, processedChunks)) {
            return this.createNullResult();
        }
        return this.createResult(typeAlias);
    }
}