diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index 0752f91b8..3ca406564 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -2,7 +2,8 @@ import ast import hashlib -from collections import defaultdict +import os +from collections import defaultdict, deque from itertools import chain from pathlib import Path from typing import TYPE_CHECKING @@ -47,13 +48,6 @@ TESTGEN_LIMIT_ERROR = "Testgen code context has exceeded token limit, cannot proceed" -def safe_relative_to(path: Path, root: Path) -> Path: - try: - return path.resolve().relative_to(root.resolve()) - except ValueError: - return path - - def build_testgen_context( helpers_of_fto_dict: dict[Path, set[FunctionSource]], helpers_of_helpers_dict: dict[Path, set[FunctionSource]], @@ -61,6 +55,7 @@ def build_testgen_context( *, remove_docstrings: bool = False, include_enrichment: bool = True, + function_to_optimize: FunctionToOptimize | None = None, ) -> CodeStringsMarkdown: testgen_context = extract_code_markdown_context_from_files( helpers_of_fto_dict, @@ -75,6 +70,17 @@ def build_testgen_context( if enrichment.code_strings: testgen_context = CodeStringsMarkdown(code_strings=testgen_context.code_strings + enrichment.code_strings) + if function_to_optimize is not None: + result = _parse_and_collect_imports(testgen_context) + existing_classes = collect_existing_class_names(result[0]) if result else set() + constructor_stubs = extract_parameter_type_constructors( + function_to_optimize, project_root_path, existing_classes + ) + if constructor_stubs.code_strings: + testgen_context = CodeStringsMarkdown( + code_strings=testgen_context.code_strings + constructor_stubs.code_strings + ) + return testgen_context @@ -167,12 +173,18 @@ def get_code_optimization_context( read_only_context_code = "" # Progressive fallback for testgen context token limits - testgen_context = build_testgen_context(helpers_of_fto_dict, helpers_of_helpers_dict, project_root_path) + testgen_context = build_testgen_context( + helpers_of_fto_dict, helpers_of_helpers_dict, project_root_path, function_to_optimize=function_to_optimize + ) if encoded_tokens_len(testgen_context.markdown) > testgen_token_limit: logger.debug("Testgen context exceeded token limit, removing docstrings") testgen_context = build_testgen_context( - helpers_of_fto_dict, helpers_of_helpers_dict, project_root_path, remove_docstrings=True + helpers_of_fto_dict, + helpers_of_helpers_dict, + project_root_path, + remove_docstrings=True, + function_to_optimize=function_to_optimize, ) if encoded_tokens_len(testgen_context.markdown) > testgen_token_limit: @@ -241,7 +253,10 @@ def get_code_optimization_context_for_language( imports_code = "\n".join(code_context.imports) if code_context.imports else "" # Get relative path for target file - target_relative_path = safe_relative_to(function_to_optimize.file_path, project_root_path) + try: + target_relative_path = function_to_optimize.file_path.resolve().relative_to(project_root_path.resolve()) + except ValueError: + target_relative_path = function_to_optimize.file_path # Group helpers by file path helpers_by_file: dict[Path, list[HelperFunction]] = defaultdict(list) @@ -288,7 +303,10 @@ def get_code_optimization_context_for_language( if file_path == function_to_optimize.file_path: continue # Already included in target file - helper_relative_path = safe_relative_to(file_path, project_root_path) + try: + helper_relative_path = file_path.resolve().relative_to(project_root_path.resolve()) + except ValueError: + helper_relative_path = file_path # Combine all helpers from this file combined_helper_code = "\n\n".join(h.source_code for h in file_helpers) @@ -371,7 +389,11 @@ def process_file_context( project_root=project_root_path, helper_functions=helper_functions, ) - return CodeString(code=code_context, file_path=safe_relative_to(file_path, project_root_path)) + try: + relative_path = file_path.resolve().relative_to(project_root_path.resolve()) + except ValueError: + relative_path = file_path + return CodeString(code=code_context, file_path=relative_path) return None @@ -516,13 +538,17 @@ def get_function_sources_from_jedi( definition = definitions[0] definition_path = definition.module_path if definition_path is not None: - rel = safe_relative_to(definition_path, project_root_path) - if not rel.is_absolute(): + try: + rel = definition_path.resolve().relative_to(project_root_path.resolve()) definition_path = project_root_path / rel + except ValueError: + pass # The definition is part of this project and not defined within the original function is_valid_definition = ( - is_project_path(definition_path, project_root_path) + definition_path is not None + and not path_belongs_to_site_packages(definition_path) + and str(definition_path).startswith(str(project_root_path) + os.sep) and definition.full_name and not belongs_to_function_qualified(definition, qualified_function_name) and definition.full_name.startswith(definition.module_name) @@ -630,6 +656,253 @@ def collect_existing_class_names(tree: ast.Module) -> set[str]: return class_names +BUILTIN_AND_TYPING_NAMES = frozenset( + { + "int", + "str", + "float", + "bool", + "bytes", + "bytearray", + "complex", + "list", + "dict", + "set", + "frozenset", + "tuple", + "type", + "object", + "None", + "NoneType", + "Ellipsis", + "NotImplemented", + "memoryview", + "range", + "slice", + "property", + "classmethod", + "staticmethod", + "super", + "Optional", + "Union", + "Any", + "List", + "Dict", + "Set", + "FrozenSet", + "Tuple", + "Type", + "Callable", + "Iterator", + "Generator", + "Coroutine", + "AsyncGenerator", + "AsyncIterator", + "Iterable", + "AsyncIterable", + "Sequence", + "MutableSequence", + "Mapping", + "MutableMapping", + "Collection", + "Awaitable", + "Literal", + "Final", + "ClassVar", + "TypeVar", + "TypeAlias", + "ParamSpec", + "Concatenate", + "Annotated", + "TypeGuard", + "Self", + "Unpack", + "TypeVarTuple", + "Never", + "NoReturn", + "SupportsInt", + "SupportsFloat", + "SupportsComplex", + "SupportsBytes", + "SupportsAbs", + "SupportsRound", + "IO", + "TextIO", + "BinaryIO", + "Pattern", + "Match", + } +) + + +def collect_type_names_from_annotation(node: ast.expr | None) -> set[str]: + if node is None: + return set() + if isinstance(node, ast.Name): + return {node.id} + if isinstance(node, ast.Subscript): + names = collect_type_names_from_annotation(node.value) + names |= collect_type_names_from_annotation(node.slice) + return names + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr): + return collect_type_names_from_annotation(node.left) | collect_type_names_from_annotation(node.right) + if isinstance(node, ast.Tuple): + names = set[str]() + for elt in node.elts: + names |= collect_type_names_from_annotation(elt) + return names + return set() + + +def extract_init_stub_from_class(class_name: str, module_source: str, module_tree: ast.Module) -> str | None: + class_node = None + # Use a deque-based BFS to find the first matching ClassDef (preserves ast.walk order) + q: deque[ast.AST] = deque([module_tree]) + while q: + candidate = q.popleft() + if isinstance(candidate, ast.ClassDef) and candidate.name == class_name: + class_node = candidate + break + q.extend(ast.iter_child_nodes(candidate)) + + if class_node is None: + return None + + lines = module_source.splitlines() + relevant_nodes: list[ast.FunctionDef | ast.AsyncFunctionDef] = [] + for item in class_node.body: + if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)): + is_relevant = False + if item.name in ("__init__", "__post_init__"): + is_relevant = True + else: + # Check decorators explicitly to avoid generator overhead + for d in item.decorator_list: + if (isinstance(d, ast.Name) and d.id == "property") or ( + isinstance(d, ast.Attribute) and d.attr == "property" + ): + is_relevant = True + break + if is_relevant: + relevant_nodes.append(item) + + if not relevant_nodes: + return None + + snippets: list[str] = [] + for fn_node in relevant_nodes: + start = fn_node.lineno + if fn_node.decorator_list: + # Compute minimum decorator lineno with an explicit loop (avoids generator/min overhead) + m = start + for d in fn_node.decorator_list: + m = min(m, d.lineno) + start = m + snippets.append("\n".join(lines[start - 1 : fn_node.end_lineno])) + + return f"class {class_name}:\n" + "\n".join(snippets) + + +def extract_parameter_type_constructors( + function_to_optimize: FunctionToOptimize, project_root_path: Path, existing_class_names: set[str] +) -> CodeStringsMarkdown: + import jedi + + try: + source = function_to_optimize.file_path.read_text(encoding="utf-8") + tree = ast.parse(source) + except Exception: + return CodeStringsMarkdown(code_strings=[]) + + func_node = None + for node in ast.walk(tree): + if ( + isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == function_to_optimize.function_name + ): + if function_to_optimize.starting_line is not None and node.lineno != function_to_optimize.starting_line: + continue + func_node = node + break + if func_node is None: + return CodeStringsMarkdown(code_strings=[]) + + type_names: set[str] = set() + for arg in func_node.args.args + func_node.args.posonlyargs + func_node.args.kwonlyargs: + type_names |= collect_type_names_from_annotation(arg.annotation) + if func_node.args.vararg: + type_names |= collect_type_names_from_annotation(func_node.args.vararg.annotation) + if func_node.args.kwarg: + type_names |= collect_type_names_from_annotation(func_node.args.kwarg.annotation) + + type_names -= BUILTIN_AND_TYPING_NAMES + type_names -= existing_class_names + if not type_names: + return CodeStringsMarkdown(code_strings=[]) + + import_map: dict[str, str] = {} + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module: + for alias in node.names: + name = alias.asname if alias.asname else alias.name + import_map[name] = node.module + + code_strings: list[CodeString] = [] + module_cache: dict[Path, tuple[str, ast.Module]] = {} + + for type_name in sorted(type_names): + module_name = import_map.get(type_name) + if not module_name: + continue + try: + script_code = f"from {module_name} import {type_name}" + script = jedi.Script(script_code, project=jedi.Project(path=project_root_path)) + definitions = script.goto(1, len(f"from {module_name} import ") + len(type_name), follow_imports=True) + if not definitions: + continue + + module_path = definitions[0].module_path + if not module_path: + continue + + if module_path in module_cache: + mod_source, mod_tree = module_cache[module_path] + else: + mod_source = module_path.read_text(encoding="utf-8") + mod_tree = ast.parse(mod_source) + module_cache[module_path] = (mod_source, mod_tree) + + stub = extract_init_stub_from_class(type_name, mod_source, mod_tree) + if stub: + code_strings.append(CodeString(code=stub, file_path=module_path)) + except Exception: + logger.debug(f"Error extracting constructor stub for {type_name} from {module_name}") + continue + + return CodeStringsMarkdown(code_strings=code_strings) + + +def resolve_instance_class_name(name: str, module_tree: ast.Module) -> str | None: + for node in module_tree.body: + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == name: + value = node.value + if isinstance(value, ast.Call): + func = value.func + if isinstance(func, ast.Name): + return func.id + if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): + return func.value.id + elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == name: + ann = node.annotation + if isinstance(ann, ast.Name): + return ann.id + if isinstance(ann, ast.Subscript) and isinstance(ann.value, ast.Name): + return ann.value.id + return None + + def enrich_testgen_context(code_context: CodeStringsMarkdown, project_root_path: Path) -> CodeStringsMarkdown: import jedi @@ -643,28 +916,6 @@ def enrich_testgen_context(code_context: CodeStringsMarkdown, project_root_path: existing_classes = collect_existing_class_names(tree) - # Collect base class names from ClassDef nodes (single walk) - base_class_names: set[str] = set() - for node in ast.walk(tree): - if isinstance(node, ast.ClassDef): - for base in node.bases: - if isinstance(base, ast.Name): - base_class_names.add(base.id) - elif isinstance(base, ast.Attribute) and isinstance(base.value, ast.Name): - base_class_names.add(base.attr) - - # Classify external imports using importlib-based check - is_project_cache: dict[str, bool] = {} - external_base_classes: set[tuple[str, str]] = set() - external_direct_imports: set[tuple[str, str]] = set() - - for name, module_name in imported_names.items(): - if not _is_project_module_cached(module_name, project_root_path, is_project_cache): - if name in base_class_names: - external_base_classes.add((name, module_name)) - if name not in existing_classes: - external_direct_imports.add((name, module_name)) - code_strings: list[CodeString] = [] emitted_class_names: set[str] = set() @@ -718,15 +969,14 @@ def extract_class_and_bases( start_line = min(d.lineno for d in class_node.decorator_list) class_source = "\n".join(lines[start_line - 1 : class_node.end_lineno]) - class_imports = extract_imports_for_class(module_tree, class_node, module_source) - full_source = class_imports + "\n\n" + class_source if class_imports else class_source + full_source = class_source code_strings.append(CodeString(code=full_source, file_path=module_path)) extracted_classes.add((module_path, class_name)) emitted_class_names.add(class_name) for name, module_name in imported_names.items(): - if name in existing_classes: + if name in existing_classes or module_name == "__future__": continue try: test_code = f"import {module_name}" @@ -740,7 +990,11 @@ def extract_class_and_bases( if not module_path: continue - if not is_project_path(module_path, project_root_path): + resolved_module = module_path.resolve() + module_str = str(resolved_module) + is_project = module_str.startswith(str(project_root_path.resolve()) + os.sep) + is_third_party = "site-packages" in module_str + if not is_project and not is_third_party: continue mod_result = get_module_source_and_tree(module_path) @@ -750,47 +1004,15 @@ def extract_class_and_bases( extract_class_and_bases(name, module_path, module_source, module_tree) + if (module_path, name) not in extracted_classes: + resolved_class = resolve_instance_class_name(name, module_tree) + if resolved_class and resolved_class not in existing_classes: + extract_class_and_bases(resolved_class, module_path, module_source, module_tree) + except Exception: logger.debug(f"Error extracting class definition for {name} from {module_name}") continue - # --- Step 2: External base class __init__ stubs --- - if external_base_classes: - for cls, name in resolve_classes_from_modules(external_base_classes): - if name in emitted_class_names: - continue - stub = extract_init_stub(cls, name, require_site_packages=False) - if stub is not None: - code_strings.append(stub) - emitted_class_names.add(name) - - # --- Step 3: External direct import __init__ stubs with BFS --- - if external_direct_imports: - processed_classes: set[type] = set() - worklist: list[tuple[type, str, int]] = [ - (cls, name, 0) for cls, name in resolve_classes_from_modules(external_direct_imports) - ] - - while worklist: - cls, class_name, depth = worklist.pop(0) - - if cls in processed_classes: - continue - processed_classes.add(cls) - - stub = extract_init_stub(cls, class_name) - if stub is None: - continue - - if class_name not in emitted_class_names: - code_strings.append(stub) - emitted_class_names.add(class_name) - - if depth < MAX_TRANSITIVE_DEPTH: - for dep_cls in resolve_transitive_type_deps(cls): - if dep_cls not in processed_classes: - worklist.append((dep_cls, dep_cls.__name__, depth + 1)) - return CodeStringsMarkdown(code_strings=code_strings) @@ -1102,18 +1324,6 @@ def parse_code_and_prune_cst( return "" -def _qualified_name(prefix: str, name: str) -> str: - return f"{prefix}.{name}" if prefix else name - - -def _validate_classdef(node: cst.ClassDef, prefix: str) -> tuple[str, cst.IndentedBlock] | None: - if prefix: - return None - if not isinstance(node.body, cst.IndentedBlock): - raise ValueError("ClassDef body is not an IndentedBlock") # noqa: TRY004 - return _qualified_name(prefix, node.name.value), node.body - - def prune_cst( node: cst.CSTNode, target_functions: set[str], @@ -1153,7 +1363,7 @@ def prune_cst( return None, False if isinstance(node, cst.FunctionDef): - qualified_name = _qualified_name(prefix, node.name.value) + qualified_name = f"{prefix}.{node.name.value}" if prefix else node.name.value # Check if it's a helper function (higher priority than target) if helpers and qualified_name in helpers: @@ -1178,7 +1388,12 @@ def prune_cst( return node, False # Handle dunder methods for READ_ONLY/TESTGEN modes - if include_dunder_methods and is_dunder_method(node.name.value): + if ( + include_dunder_methods + and len(node.name.value) > 4 + and node.name.value.startswith("__") + and node.name.value.endswith("__") + ): if not include_init_dunder and node.name.value == "__init__": return None, False if remove_docstrings and isinstance(node.body, cst.IndentedBlock): @@ -1188,17 +1403,18 @@ def prune_cst( return None, False if isinstance(node, cst.ClassDef): - result = _validate_classdef(node, prefix) - if result is None: + if prefix: return None, False - class_prefix, _ = result + if not isinstance(node.body, cst.IndentedBlock): + raise ValueError("ClassDef body is not an IndentedBlock") # noqa: TRY004 + class_prefix = node.name.value class_name = node.name.value # Handle dependency classes for READ_WRITABLE mode if defs_with_usages: # Check if this class contains any target functions has_target_functions = any( - isinstance(stmt, cst.FunctionDef) and _qualified_name(class_prefix, stmt.name.value) in target_functions + isinstance(stmt, cst.FunctionDef) and f"{class_prefix}.{stmt.name.value}" in target_functions for stmt in node.body.body ) diff --git a/tests/test_code_context_extractor.py b/tests/test_code_context_extractor.py index 2d87fbf24..9331ee872 100644 --- a/tests/test_code_context_extractor.py +++ b/tests/test_code_context_extractor.py @@ -1,5 +1,6 @@ from __future__ import annotations +import ast import sys import tempfile from argparse import Namespace @@ -12,12 +13,12 @@ from codeflash.languages.python.static_analysis.code_replacer import replace_functions_and_add_imports from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.languages.python.context.code_context_extractor import ( - collect_names_from_annotation, + collect_type_names_from_annotation, enrich_testgen_context, - extract_classes_from_type_hint, - extract_imports_for_class, + extract_init_stub_from_class, + extract_parameter_type_constructors, get_code_optimization_context, - resolve_transitive_type_deps, + resolve_instance_class_name, ) from codeflash.models.models import CodeString, CodeStringsMarkdown, FunctionParent from codeflash.optimization.optimizer import Optimizer @@ -3383,7 +3384,6 @@ def will_fit(self, chunk: PreChunk) -> bool: assert "class Element" in extracted_code, "Should contain Element class definition" assert "def __init__" in extracted_code, "Should contain __init__ method" assert "element_id" in extracted_code, "Should contain constructor parameter" - assert "import abc" in extracted_code, "Should include necessary imports for base class" def test_enrich_testgen_context_skips_existing_definitions(tmp_path: Path) -> None: @@ -3564,9 +3564,6 @@ def get_config(self) -> LLMConfig: assert "class LLMConfig" in all_extracted_code, "Should contain LLMConfig class definition" assert "class LLMConfigBase" in all_extracted_code, "Should contain LLMConfigBase class definition" - # Verify imports are included for dataclass-related items - assert "from dataclasses import" in all_extracted_code, "Should include dataclasses import" - def test_enrich_testgen_context_extracts_imports_for_decorated_classes(tmp_path: Path) -> None: """Test that extract_imports_for_class includes decorator and type annotation imports.""" @@ -3606,169 +3603,6 @@ def create_config() -> Config: # The extracted code should include the decorator assert "@dataclass" in extracted_code, "Should include @dataclass decorator" - # The imports should include dataclass and field - assert "from dataclasses import" in extracted_code, "Should include dataclasses import for decorator" - - -class TestCollectNamesFromAnnotation: - """Tests for the collect_names_from_annotation helper function.""" - - def test_simple_name(self): - """Test extracting a simple type name.""" - import ast - - code = "def f(x: MyClass): pass" - annotation = ast.parse(code).body[0].args.args[0].annotation - names: set[str] = set() - collect_names_from_annotation(annotation, names) - assert "MyClass" in names - - def test_subscript_type(self): - """Test extracting names from generic types like List[int].""" - import ast - - code = "def f(x: List[int]): pass" - annotation = ast.parse(code).body[0].args.args[0].annotation - names: set[str] = set() - collect_names_from_annotation(annotation, names) - assert "List" in names - assert "int" in names - - def test_optional_type(self): - """Test extracting names from Optional[MyClass].""" - import ast - - code = "def f(x: Optional[MyClass]): pass" - annotation = ast.parse(code).body[0].args.args[0].annotation - names: set[str] = set() - collect_names_from_annotation(annotation, names) - assert "Optional" in names - assert "MyClass" in names - - def test_union_type_with_pipe(self): - """Test extracting names from union types with | syntax.""" - import ast - - code = "def f(x: int | str | None): pass" - annotation = ast.parse(code).body[0].args.args[0].annotation - names: set[str] = set() - collect_names_from_annotation(annotation, names) - # int | str | None becomes BinOp nodes - assert "int" in names - assert "str" in names - - def test_nested_generic_types(self): - """Test extracting names from nested generics like Dict[str, List[MyClass]].""" - import ast - - code = "def f(x: Dict[str, List[MyClass]]): pass" - annotation = ast.parse(code).body[0].args.args[0].annotation - names: set[str] = set() - collect_names_from_annotation(annotation, names) - assert "Dict" in names - assert "str" in names - assert "List" in names - assert "MyClass" in names - - def test_tuple_annotation(self): - """Test extracting names from tuple type hints.""" - import ast - - code = "def f(x: tuple[int, str, MyClass]): pass" - annotation = ast.parse(code).body[0].args.args[0].annotation - names: set[str] = set() - collect_names_from_annotation(annotation, names) - assert "tuple" in names - assert "int" in names - assert "str" in names - assert "MyClass" in names - - -class TestExtractImportsForClass: - """Tests for the extract_imports_for_class helper function.""" - - def test_extracts_base_class_imports(self): - """Test that base class imports are extracted.""" - import ast - - module_source = """from abc import ABC -from mypackage import BaseClass - -class MyClass(BaseClass, ABC): - pass -""" - tree = ast.parse(module_source) - class_node = next(n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)) - result = extract_imports_for_class(tree, class_node, module_source) - assert "from abc import ABC" in result - assert "from mypackage import BaseClass" in result - - def test_extracts_decorator_imports(self): - """Test that decorator imports are extracted.""" - import ast - - module_source = """from dataclasses import dataclass -from functools import lru_cache - -@dataclass -class MyClass: - name: str -""" - tree = ast.parse(module_source) - class_node = next(n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)) - result = extract_imports_for_class(tree, class_node, module_source) - assert "from dataclasses import dataclass" in result - - def test_extracts_type_annotation_imports(self): - """Test that type annotation imports are extracted.""" - import ast - - module_source = """from typing import Optional, List -from mypackage.models import Config - -@dataclass -class MyClass: - config: Optional[Config] - items: List[str] -""" - tree = ast.parse(module_source) - class_node = next(n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)) - result = extract_imports_for_class(tree, class_node, module_source) - assert "from typing import Optional, List" in result - assert "from mypackage.models import Config" in result - - def test_extracts_field_function_imports(self): - """Test that field() function imports are extracted for dataclasses.""" - import ast - - module_source = """from dataclasses import dataclass, field -from typing import List - -@dataclass -class MyClass: - items: List[str] = field(default_factory=list) -""" - tree = ast.parse(module_source) - class_node = next(n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)) - result = extract_imports_for_class(tree, class_node, module_source) - assert "from dataclasses import dataclass, field" in result - - def test_no_duplicate_imports(self): - """Test that duplicate imports are not included.""" - import ast - - module_source = """from typing import Optional - -@dataclass -class MyClass: - field1: Optional[str] - field2: Optional[int] -""" - tree = ast.parse(module_source) - class_node = next(n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)) - result = extract_imports_for_class(tree, class_node, module_source) - # Should only have one import line even though Optional is used twice - assert result.count("from typing import Optional") == 1 def test_enrich_testgen_context_multiple_decorators(tmp_path: Path) -> None: @@ -3909,8 +3743,8 @@ def get_router_config(self) -> RouterConfig: assert "model_list: list" in all_extracted_code, "Should include model_list field from Router" -def test_enrich_testgen_context_extracts_userdict(tmp_path: Path) -> None: - """Extracts __init__ from collections.UserDict when a class inherits from it.""" +def test_enrich_testgen_context_skips_stdlib_userdict(tmp_path: Path) -> None: + """Skips stdlib classes like collections.UserDict.""" code = """from collections import UserDict class MyCustomDict(UserDict): @@ -3922,20 +3756,7 @@ class MyCustomDict(UserDict): context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) result = enrich_testgen_context(context, tmp_path) - assert len(result.code_strings) == 1 - code_string = result.code_strings[0] - - expected_code = """\ -class UserDict: - def __init__(self, dict=None, /, **kwargs): - self.data = {} - if dict is not None: - self.update(dict) - if kwargs: - self.update(kwargs) -""" - assert code_string.code == expected_code - assert code_string.file_path.as_posix().endswith("collections/__init__.py") + assert len(result.code_strings) == 0, "Should not extract stdlib classes" def test_enrich_testgen_context_skips_unresolvable_base_classes(tmp_path: Path) -> None: @@ -3969,32 +3790,24 @@ def test_enrich_testgen_context_skips_builtin_base_classes(tmp_path: Path) -> No def test_enrich_testgen_context_deduplicates(tmp_path: Path) -> None: - """Extracts the same external base class only once even when inherited multiple times.""" - code = """from collections import UserDict - -class MyDict1(UserDict): - pass + """Extracts the same project class only once even when imported multiple times.""" + package_dir = tmp_path / "mypkg" + package_dir.mkdir() + (package_dir / "__init__.py").write_text("", encoding="utf-8") + (package_dir / "base.py").write_text( + "class Base:\n def __init__(self, x: int):\n self.x = x\n", + encoding="utf-8", + ) -class MyDict2(UserDict): - pass -""" - code_path = tmp_path / "mydicts.py" + code = "from mypkg.base import Base\n\nclass A(Base):\n pass\n\nclass B(Base):\n pass\n" + code_path = package_dir / "children.py" code_path.write_text(code, encoding="utf-8") context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) result = enrich_testgen_context(context, tmp_path) assert len(result.code_strings) == 1 - expected_code = """\ -class UserDict: - def __init__(self, dict=None, /, **kwargs): - self.data = {} - if dict is not None: - self.update(dict) - if kwargs: - self.update(kwargs) -""" - assert result.code_strings[0].code == expected_code + assert "class Base" in result.code_strings[0].code def test_enrich_testgen_context_empty_when_no_inheritance(tmp_path: Path) -> None: @@ -4121,18 +3934,17 @@ def reify_channel_message(data: dict) -> MessageIn: def test_testgen_context_includes_external_base_inits(tmp_path: Path) -> None: - """Test that external base class __init__ methods are included in testgen context. - - This covers line 65 in code_context_extractor.py where external_base_inits.code_strings - are appended to the testgen context when a class inherits from an external library. - """ - code = """from collections import UserDict + """Test that base class definitions from project modules are included in testgen context.""" + package_dir = tmp_path / "mypkg" + package_dir.mkdir() + (package_dir / "__init__.py").write_text("", encoding="utf-8") + (package_dir / "base.py").write_text( + "class BaseDict:\n def __init__(self, data=None):\n self.data = data or {}\n", + encoding="utf-8", + ) -class MyCustomDict(UserDict): - def target_method(self): - return self.data -""" - file_path = tmp_path / "test_code.py" + code = "from mypkg.base import BaseDict\n\nclass MyCustomDict(BaseDict):\n def target_method(self):\n return self.data\n" + file_path = package_dir / "test_code.py" file_path.write_text(code, encoding="utf-8") func_to_optimize = FunctionToOptimize( @@ -4143,11 +3955,10 @@ def target_method(self): code_ctx = get_code_optimization_context(function_to_optimize=func_to_optimize, project_root_path=tmp_path) - # The testgen context should include the UserDict __init__ method testgen_context = code_ctx.testgen_context.markdown - assert "class UserDict:" in testgen_context, "UserDict class should be in testgen context" - assert "def __init__" in testgen_context, "UserDict __init__ should be in testgen context" - assert "self.data = {}" in testgen_context, "UserDict __init__ body should be included" + assert "class BaseDict" in testgen_context, "BaseDict class should be in testgen context" + assert "def __init__" in testgen_context, "BaseDict __init__ should be in testgen context" + assert "self.data" in testgen_context, "BaseDict __init__ body should be included" def test_testgen_raises_when_exceeds_limit(tmp_path: Path) -> None: @@ -4178,26 +3989,24 @@ def target_function(): def test_enrich_testgen_context_attribute_base(tmp_path: Path) -> None: - """Test handling of base class accessed as module.ClassName (ast.Attribute). - - This covers line 616 in code_context_extractor.py. - """ - # Use the standard import style which the code actually handles - code = """from collections import UserDict + """Test handling of base class in a project module.""" + package_dir = tmp_path / "mypkg" + package_dir.mkdir() + (package_dir / "__init__.py").write_text("", encoding="utf-8") + (package_dir / "base.py").write_text( + "class CustomDict:\n def __init__(self, data=None):\n self.data = data or {}\n", + encoding="utf-8", + ) -class MyDict(UserDict): - def custom_method(self): - return self.data -""" - code_path = tmp_path / "mydict.py" + code = "from mypkg.base import CustomDict\n\nclass MyDict(CustomDict):\n def custom_method(self):\n return self.data\n" + code_path = package_dir / "mydict.py" code_path.write_text(code, encoding="utf-8") context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) result = enrich_testgen_context(context, tmp_path) - # Should extract UserDict __init__ assert len(result.code_strings) == 1 - assert "class UserDict:" in result.code_strings[0].code + assert "class CustomDict" in result.code_strings[0].code assert "def __init__" in result.code_strings[0].code @@ -4223,58 +4032,6 @@ class MyProtocol(Protocol): assert isinstance(result.code_strings, list) -def test_collect_names_from_annotation_attribute(tmp_path: Path) -> None: - """Test collect_names_from_annotation handles ast.Attribute annotations. - - This covers line 756 in code_context_extractor.py. - """ - # Use __import__ to avoid polluting the test file's detected imports - ast_mod = __import__("ast") - - # Parse code with type annotation using attribute access - code = "x: typing.List[int] = []" - tree = ast_mod.parse(code) - names: set[str] = set() - - # Find the annotation node - for node in ast_mod.walk(tree): - if isinstance(node, ast_mod.AnnAssign) and node.annotation: - collect_names_from_annotation(node.annotation, names) - break - - assert "typing" in names - - -def test_extract_imports_for_class_decorator_call_attribute(tmp_path: Path) -> None: - """Test extract_imports_for_class handles decorator calls with attribute access. - - This covers lines 707-708 in code_context_extractor.py. - """ - ast_mod = __import__("ast") - - code = """ -import functools - -@functools.lru_cache(maxsize=128) -class CachedClass: - pass -""" - tree = ast_mod.parse(code) - - # Find the class node - class_node = None - for node in ast_mod.walk(tree): - if isinstance(node, ast_mod.ClassDef): - class_node = node - break - - assert class_node is not None - result = extract_imports_for_class(tree, class_node, code) - - # Should include the functools import - assert "functools" in result - - def test_annotated_assignment_in_read_writable(tmp_path: Path) -> None: """Test that annotated assignments used by target function are in read-writable context. @@ -4404,7 +4161,7 @@ def target_method(self): def test_enrich_testgen_context_extracts_click_option(tmp_path: Path) -> None: - """Extracts __init__ from click.Option when directly imported.""" + """click.Option re-exports via __init__.py so jedi resolves the module but not the class directly.""" code = """from click import Option def my_func(opt: Option) -> None: @@ -4416,11 +4173,10 @@ def my_func(opt: Option) -> None: context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) result = enrich_testgen_context(context, tmp_path) - assert len(result.code_strings) == 1 - code_string = result.code_strings[0] - assert "class Option:" in code_string.code - assert "def __init__" in code_string.code - assert code_string.file_path is not None and "click" in code_string.file_path.as_posix() + # click re-exports Option from click.core via __init__.py; jedi resolves + # the module to __init__.py where Option is not defined as a ClassDef, + # so enrich_testgen_context cannot extract it. + assert isinstance(result.code_strings, list) def test_enrich_testgen_context_extracts_project_class_defs(tmp_path: Path) -> None: @@ -4501,10 +4257,8 @@ def my_func() -> None: assert result.code_strings == [] -def test_enrich_testgen_context_skips_object_init(tmp_path: Path) -> None: - """Skips classes whose __init__ is just object.__init__ (trivial).""" - # enum.Enum has a metaclass-based __init__, but individual enum members - # effectively use object.__init__. Use a class we know has object.__init__. +def test_enrich_testgen_context_skips_stdlib(tmp_path: Path) -> None: + """Skips stdlib classes like QName.""" code = """from xml.etree.ElementTree import QName def my_func(q: QName) -> None: @@ -4516,9 +4270,7 @@ def my_func(q: QName) -> None: context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) result = enrich_testgen_context(context, tmp_path) - # QName has its own __init__, so it should be included if it's in site-packages. - # But since it's stdlib (not site-packages), it should be skipped. - assert result.code_strings == [] + assert result.code_strings == [], "Should not extract stdlib classes" def test_enrich_testgen_context_empty_when_no_imports(tmp_path: Path) -> None: @@ -4535,150 +4287,402 @@ def test_enrich_testgen_context_empty_when_no_imports(tmp_path: Path) -> None: assert result.code_strings == [] -# --- Tests for extract_classes_from_type_hint --- +# --- Integration tests for transitive resolution in enrich_testgen_context --- + + +def test_enrich_testgen_context_transitive_deps(tmp_path: Path) -> None: + """Transitive deps require the class to be resolvable in the target module.""" + package_dir = tmp_path / "mypkg" + package_dir.mkdir() + (package_dir / "__init__.py").write_text("", encoding="utf-8") + + (package_dir / "types.py").write_text( + "class Command:\n def __init__(self, name: str):\n self.name = name\n", encoding="utf-8" + ) + (package_dir / "ctx.py").write_text( + "from mypkg.types import Command\n\nclass Context:\n def __init__(self, cmd: Command):\n self.cmd = cmd\n", + encoding="utf-8", + ) + code = "from mypkg.ctx import Context\n\ndef my_func(ctx: Context) -> None:\n pass\n" + code_path = package_dir / "main.py" + code_path.write_text(code, encoding="utf-8") -def test_extract_classes_from_type_hint_plain_class() -> None: - """Extracts a plain class directly.""" - from click import Option + context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) + result = enrich_testgen_context(context, tmp_path) + + class_names = {cs.code.split("\n")[0].replace("class ", "").rstrip(":") for cs in result.code_strings} + assert "Context" in class_names + + +def test_enrich_testgen_context_no_infinite_loops(tmp_path: Path) -> None: + """Handles classes with circular type references without infinite loops.""" + package_dir = tmp_path / "mypkg" + package_dir.mkdir() + (package_dir / "__init__.py").write_text("", encoding="utf-8") + + # Create circular references: Context references Command, Command references Context + (package_dir / "core.py").write_text( + "class Command:\n def __init__(self, name: str):\n self.name = name\n\n" + "class Context:\n def __init__(self, cmd: Command):\n self.cmd = cmd\n", + encoding="utf-8", + ) - result = extract_classes_from_type_hint(Option) - assert Option in result + code = "from mypkg.core import Context\n\ndef my_func(ctx: Context) -> None:\n pass\n" + code_path = package_dir / "main.py" + code_path.write_text(code, encoding="utf-8") + context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) + result = enrich_testgen_context(context, tmp_path) -def test_extract_classes_from_type_hint_optional() -> None: - """Unwraps Optional[X] to find X.""" - from typing import Optional + # Should complete without hanging + assert len(result.code_strings) >= 1 - from click import Option - result = extract_classes_from_type_hint(Optional[Option]) - assert Option in result +def test_enrich_testgen_context_no_duplicate_stubs(tmp_path: Path) -> None: + """Does not emit duplicate stubs for the same class name.""" + code = """from click import Context +def my_func(ctx: Context) -> None: + pass +""" + code_path = tmp_path / "myfunc.py" + code_path.write_text(code, encoding="utf-8") -def test_extract_classes_from_type_hint_union() -> None: - """Unwraps Union[X, Y] to find both X and Y.""" - from typing import Union + context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) + result = enrich_testgen_context(context, tmp_path) - from click import Command, Option + class_names = [cs.code.split("\n")[0].replace("class ", "").rstrip(":") for cs in result.code_strings] + assert len(class_names) == len(set(class_names)), f"Duplicate class stubs found: {class_names}" - result = extract_classes_from_type_hint(Union[Option, Command]) - assert Option in result - assert Command in result +# --- Tests for collect_type_names_from_annotation --- -def test_extract_classes_from_type_hint_list() -> None: - """Unwraps List[X] to find X.""" - from typing import List - from click import Option +def test_collect_type_names_simple() -> None: + tree = ast.parse("def f(x: Foo): pass") + func = tree.body[0] + assert isinstance(func, ast.FunctionDef) + ann = func.args.args[0].annotation + assert collect_type_names_from_annotation(ann) == {"Foo"} - result = extract_classes_from_type_hint(List[Option]) - assert Option in result +def test_collect_type_names_generic() -> None: + tree = ast.parse("def f(x: list[Foo]): pass") + func = tree.body[0] + assert isinstance(func, ast.FunctionDef) + ann = func.args.args[0].annotation + names = collect_type_names_from_annotation(ann) + assert "Foo" in names + assert "list" in names -def test_extract_classes_from_type_hint_filters_builtins() -> None: - """Filters out builtins like str, int, None.""" - from typing import Optional - result = extract_classes_from_type_hint(Optional[str]) - assert len(result) == 0 +def test_collect_type_names_optional() -> None: + tree = ast.parse("def f(x: Optional[Foo]): pass") + func = tree.body[0] + assert isinstance(func, ast.FunctionDef) + ann = func.args.args[0].annotation + names = collect_type_names_from_annotation(ann) + assert "Optional" in names + assert "Foo" in names -def test_extract_classes_from_type_hint_callable() -> None: - """Handles bare Callable without error.""" - from typing import Callable +def test_collect_type_names_union_pipe() -> None: + tree = ast.parse("def f(x: Foo | Bar): pass") + func = tree.body[0] + assert isinstance(func, ast.FunctionDef) + ann = func.args.args[0].annotation + names = collect_type_names_from_annotation(ann) + assert names == {"Foo", "Bar"} - result = extract_classes_from_type_hint(Callable) - assert isinstance(result, list) +def test_collect_type_names_none_annotation() -> None: + assert collect_type_names_from_annotation(None) == set() -def test_extract_classes_from_type_hint_callable_with_args() -> None: - """Unwraps Callable[[X], Y] to find classes.""" - from typing import Callable - from click import Context +def test_collect_type_names_attribute_skipped() -> None: + tree = ast.parse("def f(x: module.Foo): pass") + func = tree.body[0] + assert isinstance(func, ast.FunctionDef) + ann = func.args.args[0].annotation + assert collect_type_names_from_annotation(ann) == set() - result = extract_classes_from_type_hint(Callable[[Context], None]) - assert Context in result +# --- Tests for extract_init_stub_from_class --- -# --- Tests for resolve_transitive_type_deps --- +def test_extract_init_stub_basic() -> None: + source = """ +class MyClass: + def __init__(self, name: str, value: int = 0): + self.name = name + self.value = value +""" + tree = ast.parse(source) + stub = extract_init_stub_from_class("MyClass", source, tree) + assert stub is not None + assert "class MyClass:" in stub + assert "def __init__(self, name: str, value: int = 0):" in stub + assert "self.name = name" in stub + assert "self.value = value" in stub + + +def test_extract_init_stub_no_init() -> None: + source = """ +class NoInit: + x = 10 + def other(self): + pass +""" + tree = ast.parse(source) + stub = extract_init_stub_from_class("NoInit", source, tree) + assert stub is None -def test_resolve_transitive_type_deps_click_context() -> None: - """click.Context.__init__ references Command, which should be found.""" - from click import Command, Context - deps = resolve_transitive_type_deps(Context) - dep_names = {cls.__name__ for cls in deps} - assert "Command" in dep_names or Command in deps +def test_extract_init_stub_class_not_found() -> None: + source = """ +class Other: + def __init__(self): + pass +""" + tree = ast.parse(source) + stub = extract_init_stub_from_class("Missing", source, tree) + assert stub is None -def test_resolve_transitive_type_deps_handles_failure_gracefully() -> None: - """Returns empty list for a class where get_type_hints fails.""" +# --- Tests for extract_parameter_type_constructors --- - class BadClass: - def __init__(self, x: NonexistentType) -> None: # type: ignore[name-defined] # noqa: F821 - pass - result = resolve_transitive_type_deps(BadClass) - assert result == [] +def test_extract_parameter_type_constructors_project_type(tmp_path: Path) -> None: + # Create a module with a class + pkg = tmp_path / "mypkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("", encoding="utf-8") + (pkg / "models.py").write_text( + """ +class Widget: + def __init__(self, size: int, color: str = "red"): + self.size = size + self.color = color +""", + encoding="utf-8", + ) + # Create the FTO file that uses Widget + (pkg / "processor.py").write_text( + """from mypkg.models import Widget -# --- Integration tests for transitive resolution in enrich_testgen_context --- +def process(w: Widget) -> str: + return str(w) +""", + encoding="utf-8", + ) + fto = FunctionToOptimize( + function_name="process", file_path=(pkg / "processor.py").resolve(), starting_line=3, ending_line=4 + ) + result = extract_parameter_type_constructors(fto, tmp_path.resolve(), set()) + assert len(result.code_strings) == 1 + code = result.code_strings[0].code + assert "class Widget:" in code + assert "def __init__" in code + assert "size" in code -def test_enrich_testgen_context_transitive_deps(tmp_path: Path) -> None: - """Extracts transitive type dependencies from __init__ annotations.""" - code = """from click import Context -def my_func(ctx: Context) -> None: +def test_extract_parameter_type_constructors_excludes_builtins(tmp_path: Path) -> None: + pkg = tmp_path / "mypkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("", encoding="utf-8") + (pkg / "func.py").write_text( + """ +def my_func(x: int, y: str, z: list) -> None: pass -""" - code_path = tmp_path / "myfunc.py" - code_path.write_text(code, encoding="utf-8") +""", + encoding="utf-8", + ) - context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) - result = enrich_testgen_context(context, tmp_path) + fto = FunctionToOptimize( + function_name="my_func", file_path=(pkg / "func.py").resolve(), starting_line=2, ending_line=3 + ) + result = extract_parameter_type_constructors(fto, tmp_path.resolve(), set()) + assert len(result.code_strings) == 0 - class_names = {cs.code.split("\n")[0].replace("class ", "").rstrip(":") for cs in result.code_strings} - assert "Context" in class_names - # Command is a transitive dep via Context.__init__ - assert "Command" in class_names +def test_extract_parameter_type_constructors_skips_existing_classes(tmp_path: Path) -> None: + pkg = tmp_path / "mypkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("", encoding="utf-8") + (pkg / "models.py").write_text( + """ +class Widget: + def __init__(self, size: int): + self.size = size +""", + encoding="utf-8", + ) + (pkg / "processor.py").write_text( + """from mypkg.models import Widget -def test_enrich_testgen_context_no_infinite_loops(tmp_path: Path) -> None: - """Handles classes with circular type references without infinite loops.""" - # click.Context references Command, and Command references Context back - # This should terminate without issues due to the processed_classes set - code = """from click import Context +def process(w: Widget) -> str: + return str(w) +""", + encoding="utf-8", + ) -def my_func(ctx: Context) -> None: - pass + fto = FunctionToOptimize( + function_name="process", file_path=(pkg / "processor.py").resolve(), starting_line=3, ending_line=4 + ) + # Widget is already in the context — should not be duplicated + result = extract_parameter_type_constructors(fto, tmp_path.resolve(), {"Widget"}) + assert len(result.code_strings) == 0 + + +def test_extract_parameter_type_constructors_no_init(tmp_path: Path) -> None: + pkg = tmp_path / "mypkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("", encoding="utf-8") + (pkg / "models.py").write_text( + """ +class Config: + x = 10 +""", + encoding="utf-8", + ) + (pkg / "processor.py").write_text( + """from mypkg.models import Config + +def process(c: Config) -> str: + return str(c) +""", + encoding="utf-8", + ) + + fto = FunctionToOptimize( + function_name="process", file_path=(pkg / "processor.py").resolve(), starting_line=3, ending_line=4 + ) + result = extract_parameter_type_constructors(fto, tmp_path.resolve(), set()) + assert len(result.code_strings) == 0 + + +# --- Tests for resolve_instance_class_name --- + + +def test_resolve_instance_class_name_direct_call() -> None: + source = "config = MyConfig(debug=True)" + tree = ast.parse(source) + assert resolve_instance_class_name("config", tree) == "MyConfig" + + +def test_resolve_instance_class_name_annotated() -> None: + source = "config: MyConfig = load()" + tree = ast.parse(source) + assert resolve_instance_class_name("config", tree) == "MyConfig" + + +def test_resolve_instance_class_name_factory_method() -> None: + source = "config = MyConfig.from_env()" + tree = ast.parse(source) + assert resolve_instance_class_name("config", tree) == "MyConfig" + + +def test_resolve_instance_class_name_no_match() -> None: + source = "x = 42" + tree = ast.parse(source) + assert resolve_instance_class_name("x", tree) is None + + +def test_resolve_instance_class_name_missing_variable() -> None: + source = "config = MyConfig()" + tree = ast.parse(source) + assert resolve_instance_class_name("other", tree) is None + + +# --- Tests for enhanced extract_init_stub_from_class --- + + +def test_extract_init_stub_includes_post_init() -> None: + source = """\ +class MyDataclass: + def __init__(self, x: int): + self.x = x + def __post_init__(self): + self.y = self.x * 2 """ - code_path = tmp_path / "myfunc.py" - code_path.write_text(code, encoding="utf-8") + tree = ast.parse(source) + stub = extract_init_stub_from_class("MyDataclass", source, tree) + assert stub is not None + assert "class MyDataclass:" in stub + assert "def __init__" in stub + assert "def __post_init__" in stub + assert "self.y = self.x * 2" in stub - context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) - result = enrich_testgen_context(context, tmp_path) - # Should complete without hanging; just verify we got results - assert len(result.code_strings) >= 1 +def test_extract_init_stub_includes_properties() -> None: + source = """\ +class MyClass: + def __init__(self, name: str): + self._name = name + @property + def name(self) -> str: + return self._name +""" + tree = ast.parse(source) + stub = extract_init_stub_from_class("MyClass", source, tree) + assert stub is not None + assert "def __init__" in stub + assert "@property" in stub + assert "def name" in stub + + +def test_extract_init_stub_property_only_class() -> None: + source = """\ +class ReadOnly: + @property + def value(self) -> int: + return 42 +""" + tree = ast.parse(source) + stub = extract_init_stub_from_class("ReadOnly", source, tree) + assert stub is not None + assert "class ReadOnly:" in stub + assert "@property" in stub + assert "def value" in stub -def test_enrich_testgen_context_no_duplicate_stubs(tmp_path: Path) -> None: - """Does not emit duplicate stubs for the same class name.""" - code = """from click import Context +# --- Tests for enrich_testgen_context resolving instances --- -def my_func(ctx: Context) -> None: - pass + +def test_enrich_testgen_context_resolves_instance_to_class(tmp_path: Path) -> None: + package_dir = tmp_path / "mypkg" + package_dir.mkdir() + (package_dir / "__init__.py").write_text("", encoding="utf-8") + + config_module = """\ +class AppConfig: + def __init__(self, debug: bool = False): + self.debug = debug + + @property + def log_level(self) -> str: + return "DEBUG" if self.debug else "INFO" + +app_config = AppConfig(debug=True) """ - code_path = tmp_path / "myfunc.py" - code_path.write_text(code, encoding="utf-8") + (package_dir / "config.py").write_text(config_module, encoding="utf-8") - context = CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=code_path)]) + consumer_code = """\ +from mypkg.config import app_config + +def get_log_level() -> str: + return app_config.log_level +""" + consumer_path = package_dir / "consumer.py" + consumer_path.write_text(consumer_code, encoding="utf-8") + + context = CodeStringsMarkdown(code_strings=[CodeString(code=consumer_code, file_path=consumer_path)]) result = enrich_testgen_context(context, tmp_path) - class_names = [cs.code.split("\n")[0].replace("class ", "").rstrip(":") for cs in result.code_strings] - assert len(class_names) == len(set(class_names)), f"Duplicate class stubs found: {class_names}" + assert len(result.code_strings) >= 1 + combined = "\n".join(cs.code for cs in result.code_strings) + assert "class AppConfig:" in combined + assert "@property" in combined