src/services/docstring_pipe.py - master - Watashicuvu/agentic-tools

src/services/docstring_pipe.py
686 строк31 KB
Якуб
fix(docstring): улучшить описания встроенных типов в type_semantics
04 апр 2026, 23:01
04 апр 2026, 23:01b0cab89
import asyncio
import logging
import libcst as cst
import libcst.matchers as m
import libcst.metadata as meta
from pathlib import Path
from typing import Optional, Dict, List, Set, Any
import re
from typing import TYPE_CHECKING

from src.services.heuristics_loader import HeuristicsLoader

if TYPE_CHECKING:
    from src.services.repository_context import RepositoryContext
    from src.services.docstring_llm_generator import DocstringLLMGenerator

class _DocstringApplier(cst.CSTTransformer):
    """Безопасно вставляет докстринг в конкретную функцию."""
    
    def __init__(self, target_function: str, new_docstring: str, heuristics: Dict):
        self.target_function = target_function
        self.new_docstring = new_docstring
        self.heuristics = heuristics
        self.applied = False
        self.function_context = None
        self.line_count = 0
    
    def leave_FunctionDef(
        self, 
        original_node: cst.FunctionDef, 
        updated_node: cst.FunctionDef
    ) -> cst.FunctionDef:
        # Проверяем совпадение имени
        if original_node.name.value != self.target_function:
            return updated_node
        
        self.applied = True
        self.line_count = len(cst.Module([]).code_for_node(original_node).splitlines())
        
        # Сохраняем контекст для валидации
        extractor = CodeContextExtractor(updated_node, heuristics=self.heuristics)
        self.function_context = extractor.extract_for_function(original_node)
        
        # Формируем узел докстринги
        doc_node = cst.SimpleStatementLine(
            body=[cst.Expr(value=cst.SimpleString(value=f'"""{self.new_docstring}"""'))]
        )
        
        # Вставляем/заменяем докстринг
        body_lines = list(updated_node.body.body)
        if original_node.get_docstring():
            # Заменяем существующую
            if body_lines and isinstance(body_lines[0], cst.SimpleStatementLine):
                body_lines[0] = doc_node
        else:
            # Вставляем новую
            body_lines.insert(0, doc_node)
        
        return updated_node.with_changes(
            body=updated_node.body.with_changes(body=body_lines)
        )


class CodeContextExtractor:
    """
    Статический анализатор для извлечения контекста функции.
    Все эвристики загружаются из конфигурации.
    """
    def __init__(
        self, 
        module: cst.Module, 
        filepath: Optional[Path] = None,
        heuristics: Optional[Dict] = None
    ):
        self.wrapper = meta.MetadataWrapper(module)
        self.module = self.wrapper.module
        self.parent_provider = self.wrapper.resolve(meta.ParentNodeProvider)
        self.filepath = filepath
        self.heuristics = heuristics or HeuristicsLoader.DEFAULT_HEURISTICS.copy()
        
        # Валидация структуры эвристик
        self._validate_heuristics()
        
        self.imports = self._extract_imports()
        self.class_hierarchy = self._build_class_hierarchy()

    def _validate_heuristics(self):
        """Проверяет наличие обязательных секций в эвристиках."""
        required = ["architectural_zones", "import_keywords", "type_semantics", "docstring_quality"]
        for key in required:
            if key not in self.heuristics:
                raise ValueError(f"Missing required heuristic section: '{key}'")

    def extract_for_function(self, func_node: cst.FunctionDef) -> Dict[str, Any]:
        """Полный контекст для функции/метода"""
        return {
            "module_imports": self.imports,
            "class_context": self._extract_class_context(func_node),
            "type_semantics": self._extract_type_semantics(func_node),
            "outgoing_calls": self._extract_outgoing_calls(func_node),
            "architectural_zone": self._detect_architectural_zone(),
            "complexity_metrics": self._compute_complexity(func_node),
        }

    def _extract_imports(self) -> Dict[str, List[str]]:
        """Извлекает бизнес-сущности и сервисы из импортов через конфигурируемые ключевые слова."""
        imports = {
            "business_entities": [],
            "services": [],
            "utils": [],
            "external_libs": []
        }
        
        # Предварительно компилируем паттерны для быстрого поиска
        keyword_map = {
            "business_entities": self.heuristics["import_keywords"].get("business_entities", []),
            "services": self.heuristics["import_keywords"].get("services", []),
            "utils": self.heuristics["import_keywords"].get("utils", []),
        }
        
        class ImportVisitor(cst.CSTVisitor):
            def visit_Import(self, node: cst.Import) -> None:
                for alias in node.names:
                    name = alias.name.value if isinstance(alias.name, cst.Name) else str(alias.name)
                    self._categorize_import(name, imports, keyword_map)
            
            def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
                module = node.module.value if node.module else ""
                for alias in node.names:
                    if isinstance(alias.name, cst.Name):
                        full_name = f"{module}.{alias.name.value}" if module else alias.name.value
                        self._categorize_import(full_name, imports, keyword_map)
            
            def _categorize_import(self, name: str, imports: Dict, keyword_map: Dict):
                name_lower = name.lower()
                
                # Проверяем по приоритету: сначала бизнес-сущности, потом сервисы, потом утилиты
                for category in ["business_entities", "services", "utils"]:
                    keywords = keyword_map.get(category, [])
                    if any(kw in name_lower for kw in keywords):
                        imports[category].append(name)
                        return
                
                # Всё остальное — внешние библиотеки
                imports["external_libs"].append(name)
        
        visitor = ImportVisitor()
        self.module.visit(visitor)
        return imports

    def _build_class_hierarchy(self) -> Dict[str, Dict]:
        """Строит иерархию классов в модуле для контекста методов"""
        hierarchy = {}
        
        class ClassVisitor(cst.CSTVisitor):
            def visit_ClassDef(self, node: cst.ClassDef) -> None:
                bases = []
                for base in node.bases:
                    if isinstance(base.value, cst.Name):
                        bases.append(base.value.value)
                    elif isinstance(base.value, cst.Attribute):
                        parts = []
                        expr = base.value
                        while isinstance(expr, cst.Attribute):
                            parts.append(expr.attr.value)
                            expr = expr.value
                        if isinstance(expr, cst.Name):
                            parts.append(expr.value)
                        bases.append(".".join(reversed(parts)))
                
                hierarchy[node.name.value] = {
                    "bases": bases,
                    "methods": [stmt.name.value for stmt in node.body.body 
                              if isinstance(stmt, cst.FunctionDef)],
                    "line_no": node.lineno if hasattr(node, 'lineno') else None
                }
        
        visitor = ClassVisitor()
        self.module.visit(visitor)
        return hierarchy

    def _extract_class_context(self, func_node: cst.FunctionDef) -> Optional[Dict]:
        """Определяет, является ли функция методом класса, через метаданные"""
        parent = self.parent_provider.get(func_node)
        while parent and not isinstance(parent, cst.ClassDef):
            parent = self.parent_provider.get(parent)
        
        if isinstance(parent, cst.ClassDef):
            class_name = parent.name.value
            return {
                "class_name": class_name,
                "bases": self.class_hierarchy.get(class_name, {}).get("bases", []),
                "is_static": any(
                    isinstance(decorator.decorator, cst.Name) and decorator.decorator.value == "staticmethod"
                    for decorator in func_node.decorators
                ),
                "is_property": any(
                    isinstance(decorator.decorator, cst.Name) and decorator.decorator.value == "property"
                    for decorator in func_node.decorators
                ),
                "is_async": isinstance(func_node, cst.Asynchronous)
            }
        return None

    def _extract_type_semantics(self, func_node: cst.FunctionDef) -> Dict[str, str]:
        """Извлекает семантику типов из аннотаций параметров и возврата через конфигурируемые паттерны."""
        semantics = {}
        
        for param in func_node.params.params:
            if param.annotation and param.name.value != "self":
                type_str = self._get_annotation_str(param.annotation)
                semantics[f"param:{param.name.value}"] = self._describe_type(type_str)
        
        if func_node.returns:
            type_str = self._get_annotation_str(func_node.returns)
            semantics["return"] = self._describe_type(type_str)
        
        return semantics

    def _get_annotation_str(self, annotation: cst.Annotation) -> str:
        """Преобразует CST-аннотацию в строку"""
        try:
            return cst.Module([]).code_for_node(annotation.annotation).strip()
        except:
            return str(annotation)

    def _describe_type(self, type_str: str) -> str:
        """
        Добавляет бизнес-семантику к типу через конфигурируемые паттерны.
        Если паттерн не найден — возвращает нейтральное описание.
        """
        type_lower = type_str.lower()

        # Ищем совпадение в конфигурируемых паттернах
        for rule in self.heuristics.get("type_semantics", []):
            if re.search(rule["pattern"], type_lower, re.IGNORECASE):
                return rule["description"]

        # Рекурсивная обработка коллекций
        if "list[" in type_lower or "dict[" in type_lower:
            inner = re.search(r"\[(.*?)\]", type_str)
            if inner:
                inner_desc = self._describe_type(inner.group(1))
                return f"Коллекция элементов типа {inner_desc}"

        # Описания для встроенных типов
        builtin_descriptions = {
            "str": "строковое значение",
            "int": "целочисленное значение",
            "float": "число с плавающей точкой",
            "bool": "логическое значение (True/False)",
            "bytes": "байтовая строка",
            "list": "список значений",
            "dict": "словарь (key-value)",
            "set": "множество уникальных значений",
            "tuple": "кортеж значений",
            "none": "отсутствие значения (None)",
            "any": "значение любого типа",
            "optional": "необязательное значение (может быть None)",
        }
        
        # Проверяем встроенные типы
        if type_lower in builtin_descriptions:
            return builtin_descriptions[type_lower]
        
        # Union типы
        if "|" in type_str or "union" in type_lower:
            parts = [p.strip() for p in type_str.replace("Union[", "").replace("]", "").split("|")]
            descs = [self._describe_type(p) for p in parts]
            return f"одно из: {', '.join(descs)}"

        return f"тип `{type_str}`"

    def _extract_outgoing_calls(self, func_node: cst.FunctionDef) -> List[str]:
        """Извлекает исходящие вызовы функции (зависимости) без разрешения символов."""
        calls = set()
        
        class CallVisitor(cst.CSTVisitor):
            def visit_Call(self, node: cst.Call) -> None:
                # Простое извлечение имени вызова без разрешения
                if isinstance(node.func, cst.Name):
                    calls.add(node.func.value)
                elif isinstance(node.func, cst.Attribute):
                    # Для вызовов вида self.method() или obj.method()
                    if isinstance(node.func.attr, cst.Name):
                        calls.add(node.func.attr.value)
        
        visitor = CallVisitor()
        func_node.body.visit(visitor)
        return sorted(calls)

    def _detect_architectural_zone(self) -> str:
        """
        Определяет архитектурную зону по пути к файлу через конфигурируемые паттерны.
        Приоритет: первая сработавшая эвристика с максимальным весом.
        """
        if not self.filepath:
            return "unknown"
        
        path_str = str(self.filepath).lower()
        best_match = ("unknown", 0.0)
        
        for rule in self.heuristics.get("architectural_zones", []):
            if re.search(rule["pattern"], path_str, re.IGNORECASE):
                if rule.get("weight", 0.0) > best_match[1]:
                    best_match = (rule["zone"], rule.get("weight", 0.0))
        
        return best_match[0]

    def _compute_complexity(self, func_node: cst.FunctionDef) -> Dict[str, int]:
        """Простые метрики сложности для приоритизации документирования"""
        class ComplexityVisitor(cst.CSTVisitor):
            def __init__(self):
                self.nesting_depth = 0
                self.max_nesting = 0
                self.branches = 0
            
            def visit_If(self, node: cst.If) -> None:
                self.branches += 1
                self.nesting_depth += 1
                self.max_nesting = max(self.max_nesting, self.nesting_depth)
            
            def leave_If(self, node: cst.If) -> None:
                self.nesting_depth -= 1
            
            def visit_For(self, node: cst.For) -> None:
                self.nesting_depth += 1
                self.max_nesting = max(self.max_nesting, self.nesting_depth)
            
            def leave_For(self, node: cst.For) -> None:
                self.nesting_depth -= 1
            
            def visit_While(self, node: cst.While) -> None:
                self.nesting_depth += 1
                self.max_nesting = max(self.max_nesting, self.nesting_depth)
            
            def leave_While(self, node: cst.While) -> None:
                self.nesting_depth -= 1
        
        visitor = ComplexityVisitor()
        func_node.body.visit(visitor)
        
        return {
            "param_count": len(func_node.params.params) + len(func_node.params.kwonly_params),
            "max_nesting": visitor.max_nesting,
            "branch_count": visitor.branches,
            "line_count": len(cst.Module([]).code_for_node(func_node.body).splitlines())
        }
    
    class CallSiteVisitor(cst.CSTVisitor):
        """Собирает вызовы функций/методов с разрешением импортов и точными номерами строк."""
        def __init__(
            self, 
            target_symbols: Set[str], 
            module_imports: Dict[str, str],
            position_provider: Optional[cst.metadata.PositionProvider] = None
        ):
            self.target_symbols = target_symbols
            self.module_imports = module_imports
            self.position_provider = position_provider
            self.call_sites = []
            self.current_function = "global"
            self.current_class = None
        
        def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
            self.current_function = node.name.value
        
        def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:
            self.current_function = "global"
        
        def visit_ClassDef(self, node: cst.ClassDef) -> None:
            self.current_class = node.name.value
        
        def leave_ClassDef(self, original_node: cst.ClassDef) -> None:
            self.current_class = None
        
        def visit_Call(self, node: cst.Call) -> None:
            func_name = self._resolve_call_target(node.func)
            if func_name and func_name in self.target_symbols:
                # Получаем номер строки через метаданные
                line_number = "?"
                if self.position_provider and node in self.position_provider:
                    line_number = self.position_provider[node].start.line
                
                self.call_sites.append({
                    "function_name": func_name,
                    "line_number": line_number,
                    "enclosing_function": self.current_function,
                    "enclosing_class": self.current_class
                })
        
        def _resolve_call_target(self, node) -> Optional[str]:
            if isinstance(node, cst.Name):
                return self.module_imports.get(node.value, node.value)
            
            elif isinstance(node, cst.Attribute):
                base = self._resolve_base(node.value)
                attr = node.attr.value if isinstance(node.attr, cst.Name) else str(node.attr)
                
                if base and f"{base}.{attr}" in self.target_symbols:
                    return f"{base}.{attr}"
                return attr
            
            return None
        
        def _resolve_base(self, node) -> Optional[str]:
            if isinstance(node, cst.Name):
                return self.module_imports.get(node.value, node.value)
            elif isinstance(node, cst.Attribute):
                base = self._resolve_base(node.value)
                attr = node.attr.value if isinstance(node.attr, cst.Name) else str(node.attr)
                return f"{base}.{attr}" if base else attr
            return None


class DocstringQualityChecker:
    """Хевристики для решения: нужна ли нам LLM? Все пороги конфигурируемы."""
    
    @staticmethod
    def needs_docstring(node: cst.FunctionDef, heuristics: Optional[Dict] = None) -> bool:
        """
        Определяет, требуется ли генерация докстринги.
        Параметры берутся из конфига, а не из хардкода.
        """
        cfg = heuristics["docstring_quality"]
        
        docstring = node.get_docstring()
        
        # 1. Если докстринги нет вообще → требует документирования
        if docstring is None:
            return True
        
        # 2. Если слишком короткая (порог из конфига)
        if len(docstring.strip()) < cfg["min_length"]:
            return True
        
        # 3. Если много параметров, а раздела с аргументами нет
        param_count = len(node.params.params) + len(node.params.kwonly_params)
        if param_count > cfg["param_threshold"]:
            args_keywords = cfg["required_sections"]["args_keywords"]
            if not any(kw in docstring for kw in args_keywords):
                return True
        
        # 4. Если есть аннотация возврата, а раздела с возвратом нет
        if node.returns:
            returns_keywords = cfg["required_sections"]["returns_keywords"]
            if not any(kw in docstring for kw in returns_keywords):
                return True
        
        return False


class DocstringInjector(cst.CSTTransformer):
    """Модифицирует CST-дерево, вставляя докстринги с контекстом"""

    def __init__(
        self,
        context_extractor: CodeContextExtractor,
        repository_context: Optional['RepositoryContext'] = None,
        loop: Optional[asyncio.AbstractEventLoop] = None,
        llm_generator: Optional['DocstringLLMGenerator'] = None
    ):
        self.context_extractor = context_extractor
        self.repository_context = repository_context
        self.loop = loop or asyncio.get_event_loop()
        self.llm_generator = llm_generator
        self.modified = False
        self.processed_functions: Set[str] = set()
        # Извлекаем лимиты из эвристик анализатора
        self.limits = self.context_extractor.heuristics.get("docstring_generation", {
            "max_dependencies_in_context": 3,
            "max_similar_implementations": 2
        })

    def leave_FunctionDef(
        self, 
        original_node: cst.FunctionDef, 
        updated_node: cst.FunctionDef
    ) -> cst.FunctionDef:
        # Пропускаем уже обработанные функции (защита от рекурсии)
        func_name = self._get_full_function_name(original_node)
        if func_name in self.processed_functions:
            return updated_node
        self.processed_functions.add(func_name)
        
        # 1. Проверяем качество текущей докстринги с использованием конфигурируемых порогов
        if not DocstringQualityChecker.needs_docstring(original_node, self.context_extractor.heuristics):
            return updated_node
        
        # 2. Извлекаем ПОЛНЫЙ контекст для LLM
        func_context = self.context_extractor.extract_for_function(original_node)
        
        # 3. Извлекаем код функции
        func_source = cst.Module([]).code_for_node(original_node)
        existing_doc = original_node.get_docstring()
        
        # 4. Генерация докстринги
        new_doc_content = self._generate_docstring_stub(func_source, func_context, existing_doc)
        
        # 5. Вставляем докстринг в тело функции
        docstring_node = cst.SimpleString(f'"""{new_doc_content}"""')
        body = updated_node.body
        new_body_content = list(body.body)
        
        if original_node.get_docstring():
            if new_body_content and isinstance(new_body_content[0], cst.SimpleStatementLine):
                new_body_content[0] = cst.SimpleStatementLine(
                    body=[cst.Expr(value=docstring_node)]
                )
        else:
            new_body_content.insert(0, cst.SimpleStatementLine(
                body=[cst.Expr(value=docstring_node)]
            ))
        
        self.modified = True
        return updated_node.with_changes(
            body=body.with_changes(body=new_body_content)
        )

    def _get_full_function_name(self, node: cst.FunctionDef) -> str:
        """Получает полное имя функции с учётом вложенности в классы"""
        parts = []
        parent = self.context_extractor.parent_provider.get(node)
        while parent and not isinstance(parent, cst.Module):
            if isinstance(parent, cst.ClassDef):
                parts.append(parent.name.value)
            parent = self.context_extractor.parent_provider.get(parent)
        parts.append(node.name.value)
        return ".".join(reversed(parts))

    def _generate_docstring_stub(
        self,
        func_code: str,
        context: Dict[str, Any],
        existing_doc: Optional[str]
    ) -> str:
        """
        Генерация докстринги с обогащением через семантический поиск.
        Все лимиты берутся из конфига.
        
        Если доступен LLM-генератор — использует его.
        Иначе — фолбэк на шаблонную генерацию.
        """
        # Если есть LLM-генератор — используем его
        if self.llm_generator:
            try:
                # Асинхронный вызов LLM
                result = self.loop.run_until_complete(
                    self.llm_generator.generate_docstring(
                        function_name=self._get_current_function_name(),
                        function_code=func_code,
                        context=context
                    )
                )
                return result.docstring_content
            except Exception as e:
                logging.warning(
                    f"⚠️ LLM generation failed, falling back to template: {e}"
                )
                # Фолбэк на шаблонную генерацию
                return self._generate_template_docstring(func_code, context, existing_doc)
        
        # Фолбэк: шаблонная генерация (старое поведение)
        return self._generate_template_docstring(func_code, context, existing_doc)
    
    def _get_current_function_name(self) -> str:
        """Получает имя текущей обрабатываемой функции."""
        # Используем последний добавленный в processed_functions
        if self.processed_functions:
            return list(self.processed_functions)[-1]
        return "unknown"
    
    def _generate_template_docstring(
        self,
        func_code: str,
        context: Dict[str, Any],
        existing_doc: Optional[str]
    ) -> str:
        """Шаблонная генерация докстринги (фолбэк при недоступности LLM)."""
        # Базовый контекст
        lines = [
            "Processes business logic with domain context.",
            "",
            "Business context:"
        ]

        if context["architectural_zone"] != "unknown":
            lines.append(f"  Zone: {context['architectural_zone']}")

        # Лимит зависимостей из конфига
        max_deps = self.limits["max_dependencies_in_context"]
        if context["outgoing_calls"]:
            deps = context["outgoing_calls"][:max_deps]
            lines.append(f"  Dependencies: {', '.join(deps)}")

        # АСИНХРОННЫЙ ПОИСК ПОХОЖИХ ФУНКЦИЙ (блокирующий вызов)
        similar_funcs = []
        if self.repository_context and self.repository_context.semantic_index:
            query_parts = ["implement"]
            if context.get("type_semantics"):
                params = [v for k, v in context["type_semantics"].items() if k.startswith("param:")]
                if params:
                    query_parts.append("  ".join(params[:2]))
            query = "  ".join(query_parts)

            try:
                # Лимит поиска из конфига
                max_results = self.limits["max_similar_implementations"]
                similar_funcs = self.loop.run_until_complete(
                    self.repository_context.search_implementations(
                        query=query,
                        scope=context["architectural_zone"],
                        limit=max_results
                    )
                )
            except Exception as e:
                logging.error(f"⚠️ Ошибка поиска похожих функций: {e}")

        # Добавляем результаты поиска в докстринг
        if similar_funcs:
            lines.append("")
            lines.append("Similar implementations in codebase:")
            for func in similar_funcs:
                lines.append(f"  • {func['file_path']}::{func['function_name']}  "
                           f"(relevance: {int(func['score'] * 100)}%)")

        lines.append("")
        lines.append("Args:")
        lines.append("  x: Input value (domain-specific).")
        lines.append("")
        lines.append("Returns:")
        lines.append("  Processed result with business semantics.")

        return "\n".join(lines)


def process_file(
    filepath: str,
    heuristics: Optional[Dict] = None,
    llm_generator: Optional['DocstringLLMGenerator'] = None
) -> bool:
    """
    Обрабатывает файл с полным статическим анализом контекста.
    Эвристики передаются извне для адаптации под репозиторий.
    
    Args:
        filepath: Путь к Python файлу
        heuristics: Конфигурация эвристик
        llm_generator: LLM-генератор для генерации докстрингов (опционально)
        
    Returns:
        True если файл был модифицирован
    """
    path = Path(filepath)
    if not path.exists():
        logging.warning(f"❌ File not found: {filepath}")
        return False

    try:
        original_code = path.read_text(encoding='utf-8')
        source_tree = cst.parse_module(original_code)

        # Инициализируем анализатор контекста с ЭВРИСТИКАМИ из конфига
        context_extractor = CodeContextExtractor(source_tree, path, heuristics=heuristics)

        # Запускаем трансформер с контекстом и LLM-генератором
        transformer = DocstringInjector(
            context_extractor,
            llm_generator=llm_generator
        )
        modified_tree = source_tree.visit(transformer)

        # Сохраняем только при изменениях
        if transformer.modified:
            logging.debug(f"✅ Modified {len(transformer.processed_functions)} function(s) in {filepath}")
            path.write_text(modified_tree.code, encoding='utf-8')
            return True
        else:
            logging.debug(f"⏭️ No changes needed for {filepath}")
            return False

    except Exception as e:
        logging.error(f"❌ Error processing {filepath}: {e}")
        #import traceback
        #traceback.lo.debug_exc()
        return False