gabion 0.1.0__py3-none-any.whl

gabion/refactor/engine.py

@@ -0,0 +1,726 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import libcst as cst
+
+from gabion.refactor.model import FieldSpec, RefactorPlan, RefactorRequest, TextEdit
+
+
+class RefactorEngine:
+    def __init__(self, project_root: Path | None = None) -> None:
+        self.project_root = project_root
+
+    def plan_protocol_extraction(self, request: RefactorRequest) -> RefactorPlan:
+        path = Path(request.target_path)
+        if self.project_root and not path.is_absolute():
+            path = self.project_root / path
+        try:
+            source = path.read_text()
+        except Exception as exc:
+            return RefactorPlan(errors=[f"Failed to read {path}: {exc}"])
+        try:
+            module = cst.parse_module(source)
+        except Exception as exc:
+            return RefactorPlan(errors=[f"LibCST parse failed for {path}: {exc}"])
+        protocol = (request.protocol_name or "").strip()
+        if not protocol:
+            return RefactorPlan(errors=["Protocol name is required for extraction."])
+        bundle = [name.strip() for name in request.bundle or [] if name.strip()]
+        field_specs: list[FieldSpec] = []
+        seen_fields: set[str] = set()
+        for spec in request.fields or []:
+            name = (spec.name or "").strip()
+            if not name or name in seen_fields:
+                continue
+            seen_fields.add(name)
+            field_specs.append(spec)
+        if bundle:
+            for name in bundle:
+                if name in seen_fields:
+                    continue
+                seen_fields.add(name)
+                field_specs.append(FieldSpec(name=name))
+        elif field_specs:
+            bundle = [spec.name for spec in field_specs]
+        if not bundle:
+            return RefactorPlan(errors=["Bundle fields are required for extraction."])
+
+        body = list(module.body)
+
+        insert_idx = _find_import_insert_index(body)
+
+        protocol_base: cst.CSTNode
+        import_stmt: cst.SimpleStatementLine | None = None
+        if _has_typing_import(body):
+            protocol_base = cst.Attribute(cst.Name("typing"), cst.Name("Protocol"))
+        elif _has_typing_protocol_import(body):
+            protocol_base = cst.Name("Protocol")
+        else:
+            protocol_base = cst.Name("Protocol")
+            import_stmt = cst.SimpleStatementLine(
+                [cst.ImportFrom(module=cst.Name("typing"), names=[cst.ImportAlias(name=cst.Name("Protocol"))])]
+            )
+
+        doc_lines = []
+        if request.rationale:
+            doc_lines.append(f"Rationale: {request.rationale}")
+        doc_lines.append(f"Bundle: {', '.join(bundle)}")
+        docstring = cst.SimpleStatementLine(
+            [cst.Expr(cst.SimpleString('"""' + "\\n".join(doc_lines) + '"""'))]
+        )
+        warnings: list[str] = []
+
+        def _annotation_for(hint: str | None) -> cst.BaseExpression:
+            if not hint:
+                return cst.Name("object")
+            try:
+                return cst.parse_expression(hint)
+            except Exception as exc:
+                warnings.append(f"Failed to parse type hint '{hint}': {exc}")
+                return cst.Name("object")
+
+        field_lines = []
+        for spec in field_specs:
+            field_lines.append(
+                cst.SimpleStatementLine(
+                    [
+                        cst.AnnAssign(
+                            target=cst.Name(spec.name),
+                            annotation=cst.Annotation(_annotation_for(spec.type_hint)),
+                            value=None,
+                        )
+                    ]
+                )
+            )
+        class_body = [docstring] + field_lines
+        class_def = cst.ClassDef(
+            name=cst.Name(protocol),
+            bases=[cst.Arg(protocol_base)],
+            body=cst.IndentedBlock(body=class_body),
+        )
+
+        new_body = list(body)
+        if import_stmt is not None:
+            new_body.insert(insert_idx, import_stmt)
+            insert_idx += 1
+        if insert_idx > 0 and not isinstance(new_body[insert_idx - 1], cst.EmptyLine):
+            new_body.insert(insert_idx, cst.EmptyLine())
+            insert_idx += 1
+        new_body.insert(insert_idx, class_def)
+
+        new_module = module.with_changes(body=new_body)
+
+        targets = {name.strip() for name in request.target_functions or [] if name.strip()}
+        bundle_fields = [spec.name for spec in field_specs]
+        protocol_hint = protocol
+        if targets:
+            transformer = _RefactorTransformer(
+                targets=targets,
+                bundle_fields=bundle_fields,
+                protocol_hint=protocol_hint,
+            )
+            new_module = new_module.visit(transformer)
+            warnings.extend(transformer.warnings)
+        if targets:
+            target_module = _module_name(path, self.project_root)
+            call_warnings, call_edits = _rewrite_call_sites(
+                new_module,
+                file_path=path,
+                target_path=path,
+                target_module=target_module,
+                protocol_name=protocol,
+                bundle_fields=bundle_fields,
+                targets=targets,
+            )
+            warnings.extend(call_warnings)
+            if call_edits is None:
+                new_source = new_module.code
+            else:
+                new_module = call_edits
+                new_source = new_module.code
+        else:
+            new_source = new_module.code
+        if new_source == source:
+            warnings.append("No changes generated for protocol extraction.")
+            return RefactorPlan(warnings=warnings)
+        end_line = len(source.splitlines())
+        edits = [
+            TextEdit(
+                path=str(path),
+                start=(0, 0),
+                end=(end_line, 0),
+                replacement=new_source,
+            )
+        ]
+
+        if targets and self.project_root:
+            extra_edits, extra_warnings = _rewrite_call_sites_in_project(
+                project_root=self.project_root,
+                target_path=path,
+                target_module=_module_name(path, self.project_root),
+                protocol_name=protocol,
+                bundle_fields=bundle_fields,
+                targets=targets,
+            )
+            edits.extend(extra_edits)
+            warnings.extend(extra_warnings)
+        return RefactorPlan(edits=edits, warnings=warnings)
+
+
+def _module_name(path: Path, project_root: Path | None) -> str:
+    rel = path.with_suffix("")
+    if project_root is not None:
+        try:
+            rel = rel.relative_to(project_root)
+        except ValueError:
+            pass
+    parts = list(rel.parts)
+    if parts and parts[0] == "src":
+        parts = parts[1:]
+    return ".".join(parts)
+
+
+def _is_docstring(stmt: cst.CSTNode) -> bool:
+    if not isinstance(stmt, cst.SimpleStatementLine) or not stmt.body:
+        return False
+    expr = stmt.body[0]
+    return isinstance(expr, cst.Expr) and isinstance(expr.value, cst.SimpleString)
+
+
+def _is_import(stmt: cst.CSTNode) -> bool:
+    if not isinstance(stmt, cst.SimpleStatementLine):
+        return False
+    return any(isinstance(item, (cst.Import, cst.ImportFrom)) for item in stmt.body)
+
+
+def _find_import_insert_index(body: list[cst.CSTNode]) -> int:
+    insert_idx = 0
+    if body and _is_docstring(body[0]):
+        insert_idx = 1
+    while insert_idx < len(body) and _is_import(body[insert_idx]):
+        insert_idx += 1
+    return insert_idx
+
+
+def _module_expr_to_str(expr: cst.BaseExpression | None) -> str | None:
+    if expr is None:
+        return None
+    if isinstance(expr, cst.Name):
+        return expr.value
+    if isinstance(expr, cst.Attribute):
+        parts = []
+        current: cst.BaseExpression | None = expr
+        while isinstance(current, cst.Attribute):
+            if isinstance(current.attr, cst.Name):
+                parts.append(current.attr.value)
+            current = current.value
+        if isinstance(current, cst.Name):
+            parts.append(current.value)
+        if parts:
+            return ".".join(reversed(parts))
+    return None
+
+
+def _has_typing_import(body: list[cst.CSTNode]) -> bool:
+    for stmt in body:
+        if not isinstance(stmt, cst.SimpleStatementLine):
+            continue
+        for item in stmt.body:
+            if isinstance(item, cst.Import):
+                for alias in item.names:
+                    if isinstance(alias, cst.ImportAlias) and isinstance(alias.name, cst.Name):
+                        if alias.name.value == "typing":
+                            return True
+                    if isinstance(alias, cst.ImportAlias) and isinstance(
+                        alias.name, cst.Attribute
+                    ):
+                        if _module_expr_to_str(alias.name) == "typing":
+                            return True
+    return False
+
+
+def _has_typing_protocol_import(body: list[cst.CSTNode]) -> bool:
+    for stmt in body:
+        if not isinstance(stmt, cst.SimpleStatementLine):
+            continue
+        for item in stmt.body:
+            if not isinstance(item, cst.ImportFrom):
+                continue
+            module = _module_expr_to_str(item.module)
+            if module != "typing":
+                continue
+            for alias in item.names:
+                if isinstance(alias, cst.ImportAlias) and isinstance(alias.name, cst.Name):
+                    if alias.name.value == "Protocol":
+                        return True
+    return False
+
+
+def _collect_import_context(
+    module: cst.Module,
+    *,
+    target_module: str,
+    protocol_name: str,
+) -> tuple[dict[str, str], dict[str, str], str | None]:
+    module_aliases: dict[str, str] = {}
+    imported_targets: dict[str, str] = {}
+    protocol_alias: str | None = None
+    for stmt in module.body:
+        if not isinstance(stmt, cst.SimpleStatementLine):
+            continue
+        for item in stmt.body:
+            if isinstance(item, cst.Import):
+                for alias in item.names:
+                    if not isinstance(alias, cst.ImportAlias):
+                        continue
+                    module_name = _module_expr_to_str(alias.name)
+                    if not module_name:
+                        continue
+                    if module_name != target_module:
+                        continue
+                    local = alias.asname.name.value if alias.asname else module_name
+                    module_aliases[local] = module_name
+            elif isinstance(item, cst.ImportFrom):
+                module_name = _module_expr_to_str(item.module)
+                if module_name != target_module:
+                    continue
+                for alias in item.names:
+                    if not isinstance(alias, cst.ImportAlias):
+                        continue
+                    if not isinstance(alias.name, cst.Name):
+                        continue
+                    local = alias.asname.name.value if alias.asname else alias.name.value
+                    imported_targets[local] = alias.name.value
+                    if alias.name.value == protocol_name:
+                        protocol_alias = local
+    return module_aliases, imported_targets, protocol_alias
+
+
+def _rewrite_call_sites(
+    module: cst.Module,
+    *,
+    file_path: Path,
+    target_path: Path,
+    target_module: str,
+    protocol_name: str,
+    bundle_fields: list[str],
+    targets: set[str],
+) -> tuple[list[str], cst.Module | None]:
+    warnings: list[str] = []
+    file_is_target = file_path == target_path
+    if not targets:
+        return warnings, None
+    target_simple = {name for name in targets if "." not in name}
+    target_methods: dict[str, set[str]] = {}
+    for name in targets:
+        if "." not in name:
+            continue
+        parts = name.split(".")
+        class_name = ".".join(parts[:-1])
+        method = parts[-1]
+        target_methods.setdefault(class_name, set()).add(method)
+    module_aliases: dict[str, str] = {}
+    imported_targets: dict[str, str] = {}
+    protocol_alias: str | None = None
+    if not file_is_target and target_module:
+        module_aliases, imported_targets, protocol_alias = _collect_import_context(
+            module, target_module=target_module, protocol_name=protocol_name
+        )
+
+    constructor_expr: cst.BaseExpression
+    needs_import = False
+    if file_is_target:
+        constructor_expr = cst.Name(protocol_name)
+    else:
+        if protocol_alias:
+            constructor_expr = cst.Name(protocol_alias)
+        elif module_aliases:
+            alias = sorted(module_aliases.keys())[0]
+            constructor_expr = cst.Attribute(cst.Name(alias), cst.Name(protocol_name))
+        else:
+            constructor_expr = cst.Name(protocol_name)
+            needs_import = True
+
+    transformer = _CallSiteTransformer(
+        file_is_target=file_is_target,
+        target_simple=target_simple,
+        target_methods=target_methods,
+        module_aliases=set(module_aliases.keys()),
+        imported_targets=set(
+            name for name, original in imported_targets.items() if original in target_simple
+        ),
+        bundle_fields=bundle_fields,
+        constructor_expr=constructor_expr,
+    )
+    new_module = module.visit(transformer)
+    warnings.extend(transformer.warnings)
+    if not transformer.changed:
+        return warnings, None
+
+    if not file_is_target and needs_import and target_module:
+        body = list(new_module.body)
+        insert_idx = _find_import_insert_index(body)
+        try:
+            module_expr = cst.parse_expression(target_module)
+        except Exception:
+            module_expr = cst.Name(target_module.split(".")[0])
+        if not isinstance(module_expr, (cst.Name, cst.Attribute)):
+            module_expr = cst.Name(target_module.split(".")[0])
+        import_stmt = cst.SimpleStatementLine(
+            [
+                cst.ImportFrom(
+                    module=module_expr,
+                    names=[cst.ImportAlias(name=cst.Name(protocol_name))],
+                )
+            ]
+        )
+        body.insert(insert_idx, import_stmt)
+        new_module = new_module.with_changes(body=body)
+    return warnings, new_module
+
+
+def _rewrite_call_sites_in_project(
+    *,
+    project_root: Path,
+    target_path: Path,
+    target_module: str,
+    protocol_name: str,
+    bundle_fields: list[str],
+    targets: set[str],
+) -> tuple[list[TextEdit], list[str]]:
+    edits: list[TextEdit] = []
+    warnings: list[str] = []
+    scan_root = project_root / "src"
+    if not scan_root.exists():
+        scan_root = project_root
+    for path in sorted(scan_root.rglob("*.py")):
+        if path == target_path:
+            continue
+        try:
+            source = path.read_text()
+        except Exception as exc:
+            warnings.append(f"Failed to read {path}: {exc}")
+            continue
+        try:
+            module = cst.parse_module(source)
+        except Exception as exc:
+            warnings.append(f"LibCST parse failed for {path}: {exc}")
+            continue
+        call_warnings, updated_module = _rewrite_call_sites(
+            module,
+            file_path=path,
+            target_path=target_path,
+            target_module=target_module,
+            protocol_name=protocol_name,
+            bundle_fields=bundle_fields,
+            targets=targets,
+        )
+        warnings.extend(call_warnings)
+        if updated_module is None:
+            continue
+        new_source = updated_module.code
+        if new_source == source:
+            continue
+        end_line = len(source.splitlines())
+        edits.append(
+            TextEdit(
+                path=str(path),
+                start=(0, 0),
+                end=(end_line, 0),
+                replacement=new_source,
+            )
+        )
+    return edits, warnings
+
+
+class _RefactorTransformer(cst.CSTTransformer):
+    def __init__(
+        self,
+        *,
+        targets: set[str],
+        bundle_fields: list[str],
+        protocol_hint: str,
+    ) -> None:
+        # dataflow-bundle: bundle_fields, protocol_hint, targets
+        self.targets = targets
+        self.bundle_fields = bundle_fields
+        self.protocol_hint = protocol_hint
+        self.warnings: list[str] = []
+        self._stack: list[str] = []
+
+    def visit_ClassDef(self, node: cst.ClassDef) -> bool:
+        self._stack.append(node.name.value)
+        return True
+
+    def leave_ClassDef(
+        self, original_node: cst.ClassDef, updated_node: cst.ClassDef
+    ) -> cst.CSTNode:
+        if self._stack:
+            self._stack.pop()
+        return updated_node
+
+    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
+        self._stack.append(node.name.value)
+        return True
+
+    def leave_FunctionDef(
+        self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
+    ) -> cst.CSTNode:
+        updated = self._maybe_rewrite_function(original_node, updated_node)
+        if self._stack:
+            self._stack.pop()
+        return updated
+
+    def visit_AsyncFunctionDef(self, node: cst.AsyncFunctionDef) -> bool:
+        self._stack.append(node.name.value)
+        return True
+
+    def leave_AsyncFunctionDef(
+        self, original_node: cst.AsyncFunctionDef, updated_node: cst.AsyncFunctionDef
+    ) -> cst.CSTNode:
+        updated = self._maybe_rewrite_function(original_node, updated_node)
+        if self._stack:
+            self._stack.pop()
+        return updated
+
+    def _maybe_rewrite_function(
+        self,
+        original_node: cst.FunctionDef | cst.AsyncFunctionDef,
+        updated_node: cst.FunctionDef | cst.AsyncFunctionDef,
+    ) -> cst.CSTNode:
+        qualname = ".".join(self._stack)
+        name = original_node.name.value
+        if name not in self.targets and qualname not in self.targets:
+            return updated_node
+
+        original_params = self._ordered_param_names(original_node.params)
+        if not original_params:
+            return updated_node
+
+        keep_self = original_params[0] in {"self", "cls"}
+        self_param = None
+        if keep_self:
+            self_param = self._find_self_param(original_node.params, original_params[0])
+
+        bundle_param_name = self._choose_bundle_name(original_params)
+        new_params = self._build_parameters(self_param, bundle_param_name)
+
+        bundle_set = set(self.bundle_fields)
+        target_fields = [
+            name
+            for name in original_params
+            if name in bundle_set and name not in {"self", "cls"}
+        ]
+        new_body = self._inject_preamble(
+            updated_node.body, bundle_param_name, target_fields
+        )
+        return updated_node.with_changes(params=new_params, body=new_body)
+
+    def _ordered_param_names(self, params: cst.Parameters) -> list[str]:
+        names: list[str] = []
+        for param in params.posonly_params:
+            names.append(param.name.value)
+        for param in params.params:
+            names.append(param.name.value)
+        for param in params.kwonly_params:
+            names.append(param.name.value)
+        return names
+
+    def _find_self_param(
+        self, params: cst.Parameters, name: str
+    ) -> cst.Param | None:
+        for param in params.posonly_params:
+            if param.name.value == name:
+                return param
+        for param in params.params:
+            if param.name.value == name:
+                return param
+        return None
+
+    def _choose_bundle_name(self, existing: list[str]) -> str:
+        candidate = "bundle"
+        if candidate not in existing:
+            return candidate
+        idx = 1
+        while f"bundle_{idx}" in existing:
+            idx += 1
+        return f"bundle_{idx}"
+
+    def _build_parameters(
+        self, self_param: cst.Param | None, bundle_name: str
+    ) -> cst.Parameters:
+        params: list[cst.Param] = []
+        if self_param is not None:
+            params.append(self_param)
+        annotation = None
+        if self.protocol_hint:
+            try:
+                annotation = cst.Annotation(cst.parse_expression(self.protocol_hint))
+            except Exception as exc:
+                self.warnings.append(
+                    f"Failed to parse protocol type hint '{self.protocol_hint}': {exc}"
+                )
+        params.append(
+            cst.Param(
+                name=cst.Name(bundle_name),
+                annotation=annotation,
+            )
+        )
+        return cst.Parameters(
+            params=params,
+            star_arg=cst.MaybeSentinel.DEFAULT,
+            kwonly_params=[],
+            star_kwarg=None,
+            posonly_params=[],
+            posonly_ind=cst.MaybeSentinel.DEFAULT,
+        )
+
+    def _inject_preamble(
+        self, body: cst.BaseSuite, bundle_name: str, fields: list[str]
+    ) -> cst.BaseSuite:
+        if not fields:
+            return body
+        if not isinstance(body, cst.IndentedBlock):
+            return body
+        assign_lines = [
+            cst.SimpleStatementLine(
+                [
+                    cst.Assign(
+                        targets=[cst.AssignTarget(cst.Name(name))],
+                        value=cst.Attribute(
+                            value=cst.Name(bundle_name),
+                            attr=cst.Name(name),
+                        ),
+                    )
+                ]
+            )
+            for name in fields
+        ]
+        existing = list(body.body)
+        insert_at = 0
+        if existing:
+            first = existing[0]
+            if isinstance(first, cst.SimpleStatementLine) and first.body:
+                expr = first.body[0]
+                if isinstance(expr, cst.Expr) and isinstance(expr.value, cst.SimpleString):
+                    insert_at = 1
+        new_body = existing[:insert_at] + assign_lines + existing[insert_at:]
+        return body.with_changes(body=new_body)
+
+
+class _CallSiteTransformer(cst.CSTTransformer):
+    def __init__(
+        self,
+        *,
+        file_is_target: bool,
+        target_simple: set[str],
+        target_methods: dict[str, set[str]],
+        module_aliases: set[str],
+        imported_targets: set[str],
+        bundle_fields: list[str],
+        constructor_expr: cst.BaseExpression,
+    ) -> None:
+        # dataflow-bundle: bundle_fields, constructor_expr, file_is_target, imported_targets, module_aliases, target_methods, target_simple
+        self.file_is_target = file_is_target
+        self.target_simple = target_simple
+        self.target_methods = target_methods
+        self.module_aliases = module_aliases
+        self.imported_targets = imported_targets
+        self.bundle_fields = bundle_fields
+        self.constructor_expr = constructor_expr
+        self.changed = False
+        self.warnings: list[str] = []
+        self._class_stack: list[str] = []
+
+    def visit_ClassDef(self, node: cst.ClassDef) -> bool:
+        self._class_stack.append(node.name.value)
+        return True
+
+    def leave_ClassDef(
+        self, original_node: cst.ClassDef, updated_node: cst.ClassDef
+    ) -> cst.CSTNode:
+        if self._class_stack:
+            self._class_stack.pop()
+        return updated_node
+
+    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.CSTNode:
+        if not self._is_target_call(updated_node.func):
+            return updated_node
+        if self._already_wrapped(updated_node):
+            return updated_node
+        bundle_args = self._build_bundle_args(updated_node)
+        if bundle_args is None:
+            return updated_node
+        bundle_call = cst.Call(func=self.constructor_expr, args=bundle_args)
+        self.changed = True
+        return updated_node.with_changes(args=[cst.Arg(value=bundle_call)])
+
+    def _is_target_call(self, func: cst.BaseExpression) -> bool:
+        if isinstance(func, cst.Name):
+            if self.file_is_target and func.value in self.target_simple:
+                return True
+            if not self.file_is_target and func.value in self.imported_targets:
+                return True
+            return False
+        if isinstance(func, cst.Attribute):
+            if not isinstance(func.attr, cst.Name):
+                return False
+            attr = func.attr.value
+            if self.file_is_target and self._class_stack:
+                class_name = ".".join(self._class_stack)
+                methods = self.target_methods.get(class_name, set())
+                if attr in methods and isinstance(func.value, cst.Name):
+                    if func.value.value in {"self", "cls", self._class_stack[-1]}:
+                        return True
+            if not self.file_is_target and isinstance(func.value, cst.Name):
+                if func.value.value in self.module_aliases and attr in self.target_simple:
+                    return True
+        return False
+
+    def _already_wrapped(self, call: cst.Call) -> bool:
+        if len(call.args) != 1:
+            return False
+        arg = call.args[0]
+        if arg.star:
+            return False
+        value = arg.value
+        if not isinstance(value, cst.Call):
+            return False
+        if isinstance(value.func, cst.Name) and isinstance(self.constructor_expr, cst.Name):
+            return value.func.value == self.constructor_expr.value
+        if isinstance(value.func, cst.Attribute) and isinstance(self.constructor_expr, cst.Attribute):
+            return value.func.attr.value == self.constructor_expr.attr.value
+        return False
+
+    def _build_bundle_args(self, call: cst.Call) -> list[cst.Arg] | None:
+        if any(arg.star in {"*", "**"} for arg in call.args):
+            self.warnings.append("Skipped call with star args/kwargs during refactor.")
+            return None
+        positional = [arg for arg in call.args if arg.keyword is None]
+        keyword_args = {
+            arg.keyword.value: arg.value
+            for arg in call.args
+            if arg.keyword is not None and isinstance(arg.keyword, cst.Name)
+        }
+        for key in keyword_args:
+            if key not in self.bundle_fields:
+                self.warnings.append(
+                    f"Skipped call with unknown keyword '{key}' during refactor."
+                )
+                return None
+        mapping: dict[str, cst.BaseExpression] = {}
+        mapping.update(keyword_args)
+        remaining = [field for field in self.bundle_fields if field not in mapping]
+        if len(positional) > len(remaining):
+            self.warnings.append("Skipped call with extra positional args during refactor.")
+            return None
+        for field, arg in zip(remaining, positional):
+            mapping[field] = arg.value
+        if len(mapping) != len(self.bundle_fields):
+            self.warnings.append("Skipped call with missing bundle fields during refactor.")
+            return None
+        return [
+            cst.Arg(keyword=cst.Name(field), value=mapping[field])
+            for field in self.bundle_fields
+        ]