pdd-cli 0.0.90__py3-none-any.whl → 0.0.118__py3-none-any.whl

pdd/generate_test.py

@@ -1,195 +1,256 @@
-from typing import Tuple, Optional
-from rich import print
-from rich.markdown import Markdown
+"""
+Module for generating unit tests from code or example files using LLMs.
+"""
+from __future__ import annotations
+
+import re
+from typing import Optional, Tuple
+
 from rich.console import Console
-from . import EXTRACTION_STRENGTH, DEFAULT_STRENGTH, DEFAULT_TIME
-from .load_prompt_template import load_prompt_template
-from .preprocess import preprocess
-from .llm_invoke import llm_invoke
-from .unfinished_prompt import unfinished_prompt
-from .continue_generation import continue_generation
-from .postprocess import postprocess
+from rich.markdown import Markdown
+
+from pdd import DEFAULT_STRENGTH, DEFAULT_TIME, EXTRACTION_STRENGTH
+from pdd.continue_generation import continue_generation
+from pdd.llm_invoke import llm_invoke
+from pdd.load_prompt_template import load_prompt_template
+from pdd.postprocess import postprocess
+from pdd.preprocess import preprocess
+from pdd.unfinished_prompt import unfinished_prompt
 
 console = Console()
 
+
+def _validate_inputs(
+    prompt: str,
+    code: Optional[str],
+    strength: float,
+    temperature: float,
+    language: str
+) -> None:
+    """
+    Validates the inputs for generate_test function.
+
+    Raises:
+        ValueError: If any input is invalid.
+    """
+    if not isinstance(prompt, str) or not prompt.strip():
+        raise ValueError("Prompt must be a non-empty string")
+
+    if code is None or not isinstance(code, str) or not code.strip():
+        raise ValueError("Code must be a non-empty string")
+
+    if not isinstance(strength, (int, float)) or not 0 <= strength <= 1:
+        raise ValueError("Strength must be a float between 0 and 1")
+
+    if not isinstance(temperature, (int, float)):
+        raise ValueError("Temperature must be a float")
+
+    if not isinstance(language, str) or not language.strip():
+        raise ValueError("Language must be a non-empty string")
+
+
 def generate_test(
     prompt: str,
-    code: str,
+    code: Optional[str] = None,
+    example: Optional[str] = None,
     strength: float = DEFAULT_STRENGTH,
     temperature: float = 0.0,
     time: float = DEFAULT_TIME,
-    language: str = "python",
+    language: str = 'python',
     verbose: bool = False,
     source_file_path: Optional[str] = None,
     test_file_path: Optional[str] = None,
-    module_name: Optional[str] = None
+    module_name: Optional[str] = None,
+    existing_tests: Optional[str] = None
 ) -> Tuple[str, float, str]:
     """
-    Generate a unit test from a code file using LLM.
+    Generates a unit test for a given code file or example usage using an LLM.
 
     Args:
-        prompt (str): The prompt that generated the code file.
-        code (str): The code to generate a unit test from.
-        strength (float): The strength of the LLM model (0-1).
-        temperature (float): The temperature of the LLM model.
-        language (str): The programming language for the unit test.
-        time (float, optional): Time budget for LLM calls. Defaults to DEFAULT_TIME.
-        verbose (bool): Whether to print detailed information.
-        source_file_path (Optional[str]): Absolute or relative path to the code under test.
-        test_file_path (Optional[str]): Destination path for the generated test file.
-        module_name (Optional[str]): Module name (without extension) for proper imports.
+        prompt: The prompt that generated the code (context).
+        code: The source code to test. Mutually exclusive with 'example'.
+        example: An example usage of the module. Mutually exclusive with 'code'.
+        strength: LLM strength (0.0 to 1.0).
+        temperature: LLM temperature.
+        time: Thinking effort for the LLM (0.0 to 1.0).
+        language: Target language for the test.
+        verbose: Whether to print detailed logs.
+        source_file_path: Path to the code under test.
+        test_file_path: Destination path for the test.
+        module_name: Name of the module for imports.
+        existing_tests: Content of existing tests to merge.
 
     Returns:
-        Tuple[str, float, str]: (unit_test, total_cost, model_name)
+        Tuple containing:
+        - unit_test (str): The generated test code.
+        - total_cost (float): Total cost of generation.
+        - model_name (str): Name of the model used.
     """
     total_cost = 0.0
-    model_name = ""
+    model_name = "unknown"
 
-    try:
-        # Step 1: Load prompt template
-        template = load_prompt_template("generate_test_LLM")
-        if not template:
-            raise ValueError("Failed to load generate_test_LLM prompt template")
-
-        # Step 2: Preprocess template
-        processed_template = preprocess(template, recursive=False, double_curly_brackets=False)
-        processed_prompt = preprocess(prompt, recursive=False, double_curly_brackets=False)
-
-        # Step 3: Run through LLM
-        input_json = {
-            "prompt_that_generated_code": processed_prompt,
-            "code": code,
-            "language": language,
-            "source_file_path": source_file_path or "",
-            "test_file_path": test_file_path or "",
-            "module_name": module_name or ""
-        }
+    # --- Step 1: Determine prompt template and validate inputs ---
+    if (code is None and example is None) or (code is not None and example is not None):
+        raise ValueError("Exactly one of 'code' or 'example' must be provided.")
 
-        if verbose:
-            console.print("[bold blue]Generating unit test...[/bold blue]")
+    template_name = "generate_test_from_example_LLM" if example else "generate_test_LLM"
+
+    raw_template = load_prompt_template(template_name)
+    if not raw_template:
+        raise ValueError(f"Failed to load {template_name} prompt template")
+
+    # --- Step 2: Preprocess template and prompt ---
+    # Preprocess the template
+    prompt_template = preprocess(
+        raw_template,
+        recursive=False,
+        double_curly_brackets=False
+    )
 
-        response = llm_invoke(
-            prompt=processed_template,
-            input_json=input_json,
+    # Preprocess the original prompt input
+    processed_prompt_input = preprocess(
+        prompt,
+        recursive=False,
+        double_curly_brackets=False
+    )
+
+    # --- Step 3: Run inputs through LLM ---
+    input_data = {
+        "prompt_that_generated_code": processed_prompt_input,
+        "language": language,
+        "source_file_path": source_file_path if source_file_path else "",
+        "test_file_path": test_file_path if test_file_path else "",
+        "module_name": module_name if module_name else "",
+        "existing_tests": existing_tests if existing_tests else ""
+    }
+
+    if example:
+        input_data["example"] = example
+    else:
+        input_data["code"] = code
+
+    if verbose:
+        console.print(
+            f"[bold blue]Generating unit test using template: {template_name}[/bold blue]"
+        )
+        console.print(f"[dim]Strength: {strength}, Time: {time}, Temp: {temperature}[/dim]")
+
+    try:
+        llm_result = llm_invoke(
+            prompt=prompt_template,
+            input_json=input_data,
             strength=strength,
             temperature=temperature,
             time=time,
             verbose=verbose
         )
+    except Exception as e:
+        console.print(f"[bold red]Error invoking LLM:[/bold red] {e}")
+        raise
 
-        total_cost += response['cost']
-        model_name = response['model_name']
-        result = response['result']
+    current_text = llm_result['result']
+    total_cost += llm_result.get('cost', 0.0)
+    model_name = llm_result.get('model_name', 'unknown')
 
-        # Validate that we got a non-empty result
-        if not result or not result.strip():
-            raise ValueError(f"LLM test generation returned empty result. Model: {model_name}, Cost: ${response['cost']:.6f}")
+    # --- Step 4: Verbose Output of Initial Result ---
+    if verbose:
+        console.print("[bold green]Initial LLM Output:[/bold green]")
+        console.print(Markdown(current_text))
+        console.print(f"[dim]Initial Cost: ${llm_result.get('cost', 0.0):.6f}[/dim]")
 
-        if verbose:
-            console.print(Markdown(result))
-            console.print(f"[bold green]Initial generation cost: ${total_cost:.6f}[/bold green]")
-
-        # Step 4: Check if generation is complete
-        last_600_chars = result[-600:] if len(result) > 600 else result
-        
-        # Validate that the last_600_chars is not empty after stripping
-        if not last_600_chars.strip():
-            # If the tail is empty, assume generation is complete
-            if verbose:
-                console.print("[bold yellow]Last 600 chars are empty, assuming generation is complete[/bold yellow]")
-            reasoning = "Generation appears complete (tail is empty)"
-            is_finished = True
-            check_cost = 0.0
-            check_model = model_name
-        else:
-            reasoning, is_finished, check_cost, check_model = unfinished_prompt(
-                prompt_text=last_600_chars,
-                strength=strength,
-                temperature=temperature,
-                time=time,
-                language=language,
-                verbose=verbose
-            )
-        total_cost += check_cost
-
-        if not is_finished:
-            if verbose:
-                console.print("[bold yellow]Generation incomplete. Continuing...[/bold yellow]")
-            
-            continued_result, continue_cost, continue_model = continue_generation(
-                formatted_input_prompt=processed_template,
-                llm_output=result,
+    # --- Step 5: Detect incomplete generation ---
+    # Check the last 600 characters
+    last_chunk = current_text[-600:] if len(current_text) > 600 else current_text
+
+    # Only check if there is actual content
+    if last_chunk.strip():
+        try:
+            reasoning, is_finished, check_cost, _ = unfinished_prompt(
+                prompt_text=last_chunk,
                 strength=strength,
                 temperature=temperature,
                 time=time,
                 language=language,
                 verbose=verbose
             )
-            total_cost += continue_cost
-            result = continued_result
-            model_name = continue_model
+            total_cost += check_cost
 
-        # Process the final result
-        try:
-            processed_result, post_cost, post_model = postprocess(
-                result,
-                language=language,
-                strength=EXTRACTION_STRENGTH,
-                temperature=temperature,
-                time=time,
-                verbose=verbose
-            )
-            total_cost += post_cost
+            if not is_finished:
+                if verbose:
+                    console.print(
+                        "[yellow]Output detected as incomplete. Continuing generation...[/yellow]"
+                    )
+                    console.print(f"[dim]Reasoning: {reasoning}[/dim]")
+
+                # We need the formatted prompt for continue_generation.
+                # Since llm_invoke handles formatting internally, we attempt to format here
+                # to pass context to the continuation logic.
+                try:
+                    formatted_input_prompt = prompt_template.format(**input_data)
+                except Exception:
+                    # Fallback if simple formatting fails (e.g. complex jinja or missing keys)
+                    # We use the raw template as best effort context
+                    formatted_input_prompt = prompt_template
+
+                final_llm_output, cont_cost, cont_model = continue_generation(
+                    formatted_input_prompt=formatted_input_prompt,
+                    llm_output=current_text,
+                    strength=strength,
+                    temperature=temperature,
+                    verbose=verbose
+                )
+
+                current_text = final_llm_output
+                total_cost += cont_cost
+                model_name = cont_model # Update to the model used for continuation
         except Exception as e:
-            console.print(f"[bold red]Postprocess failed: {str(e)}[/bold red]")
-            console.print(f"[bold yellow]Falling back to raw result[/bold yellow]")
-            
-            # Try to extract code blocks directly from the raw result
-            import re
-            code_blocks = re.findall(r'```(?:python)?\s*(.*?)```', result, re.DOTALL | re.IGNORECASE)
-            
-            if code_blocks:
-                # Use the first substantial code block
-                for block in code_blocks:
-                    if len(block.strip()) > 100 and ('def test_' in block or 'import' in block):
-                        processed_result = block.strip()
-                        break
-                else:
-                    processed_result = code_blocks[0].strip() if code_blocks else result
-            else:
-                # No code blocks found, use raw result
-                processed_result = result
-            
-            post_cost = 0.0
-
-        # Step 5: Print total cost if verbose
+            console.print(f"[bold red]Error during completion check/continuation:[/bold red] {e}")
+            # Proceed with what we have if check fails
+
+    # --- Step 6: Postprocess ---
+    try:
+        extracted_code, pp_cost, _ = postprocess(
+            llm_output=current_text,
+            language=language,
+            strength=EXTRACTION_STRENGTH,
+            verbose=verbose
+        )
+        total_cost += pp_cost
+        unit_test = extracted_code
+    except Exception as e:
         if verbose:
-            console.print(f"[bold green]Total cost: ${total_cost:.6f}[/bold green]")
-            console.print(f"[bold blue]Final model used: {model_name}[/bold blue]")
+            console.print(f"[bold red]Postprocessing failed:[/bold red] {e}")
+        unit_test = ""
 
-        # Step 6: Return results
-        return processed_result, total_cost, model_name
+    # Fallback extraction if postprocess returned empty or failed
+    if not unit_test.strip():
+        if verbose:
+            console.print(
+                "[yellow]Postprocess returned empty. Attempting fallback regex extraction.[/yellow]"
+            )
 
-    except Exception as e:
-        console.print(f"[bold red]Error: {str(e)}[/bold red]")
-        raise
+        # Regex to find code blocks, preferring those with specific keywords
+        code_block_pattern = re.compile(r"```(?:\w+)?\n(.*?)```", re.DOTALL)
+        matches = code_block_pattern.findall(current_text)
 
+        best_match = ""
+        for match in matches:
+            # Heuristic: prefer blocks that look like tests
+            if "def test_" in match or "import unittest" in match or "import pytest" in match:
+                best_match = match
+                break
 
-def _validate_inputs(
-    prompt: str,
-    code: str,
-    strength: float,
-    temperature: float,
-    language: str
-) -> None:
-    """Validate input parameters."""
-    if not prompt or not isinstance(prompt, str):
-        raise ValueError("Prompt must be a non-empty string")
-    if not code or not isinstance(code, str):
-        raise ValueError("Code must be a non-empty string")
-    if not isinstance(strength, float) or not 0 <= strength <= 1:
-        raise ValueError("Strength must be a float between 0 and 1")
-    if not isinstance(temperature, float):
-        raise ValueError("Temperature must be a float")
-    if not language or not isinstance(language, str):
-        raise ValueError("Language must be a non-empty string")
+        if not best_match and matches:
+            # If no specific test keywords found, take the longest block
+            best_match = max(matches, key=len)
+
+        unit_test = best_match if best_match else current_text
+
+    # --- Step 7: Final Cost Reporting ---
+    if verbose:
+        console.print(f"[bold blue]Generation Complete.[/bold blue]")
+        console.print(f"[bold]Total Cost:[/bold] ${total_cost:.6f}")
+        console.print(f"[dim]Final Model: {model_name}[/dim]")
+
+    # --- Step 8: Return ---
+    return unit_test, total_cost, model_name

pdd/get_comment.py

@@ -1,55 +1,30 @@
-# To achieve the task of writing a Python function `get_comment` that returns the comment character(s) associated with a given programming language, we need to follow the steps outlined in your description. Here's how you can implement this:
-
-# ```python
-import os
 import csv
 
-def get_comment(language):
-    # Step 1: Load environment variables to get the path to the CSV file
-    pdd_path = os.getenv('PDD_PATH')
-    if not pdd_path:
-        return 'del'  # Return 'del' if the environment variable is not set
+from pdd.path_resolution import get_default_resolver
+
+
+def get_comment(language: str) -> str:
+    try:
+        resolver = get_default_resolver()
+        csv_file_path = resolver.resolve_data_file("data/language_format.csv")
+    except ValueError:
+        return "del"
 
-    csv_file_path = os.path.join(pdd_path, 'data', 'language_format.csv')
+    if not isinstance(language, str):
+        return "del"
 
-    # Step 2: Lower case the language string
     language = language.lower()
 
     try:
-        # Step 3: Open the CSV file and look up the comment character(s)
-        with open(csv_file_path, mode='r', newline='') as csvfile:
+        with open(str(csv_file_path), mode="r", newline="") as csvfile:
             reader = csv.DictReader(csvfile)
             for row in reader:
-                if row['language'].lower() == language:
-                    comment = row['comment']
-                    # Step 4: Check if the comment character(s) is valid
-                    if comment:
-                        return comment
-                    else:
-                        return 'del'
+                if row["language"].lower() == language:
+                    comment = row.get("comment", "")
+                    return comment if comment else "del"
     except FileNotFoundError:
-        return 'del'  # Return 'del' if the file is not found
-    except Exception as e:
-        return 'del'  # Return 'del' for any other exceptions
-
-    return 'del'  # Return 'del' if the language is not found
-
-# Example usage:
-# Assuming the environment variable PDD_PATH is set correctly
-# print(get_comment('Python'))  # Output: #
-# print(get_comment('Java'))    # Output: //
-# ```
-
-# ### Explanation:
-
-# 1. **Environment Variable**: The function first retrieves the `PDD_PATH` environment variable to locate the CSV file. If the environment variable is not set, it returns `'del'`.
-
-# 2. **Case Insensitivity**: The input language string is converted to lowercase to ensure case-insensitive comparison.
-
-# 3. **CSV Reading**: The function reads the CSV file using Python's `csv.DictReader`, which allows accessing each row as a dictionary. It checks if the lowercase version of the language matches any entry in the CSV.
-
-# 4. **Validation**: If a match is found, it checks if the comment character(s) is valid (i.e., not an empty string). If valid, it returns the comment character(s); otherwise, it returns `'del'`.
-
-# 5. **Error Handling**: The function handles potential errors such as file not found or other exceptions by returning `'del'`.
+        return "del"
+    except Exception:
+        return "del"
 
-# This implementation assumes that the CSV file is correctly formatted and that the environment variable `PDD_PATH` is set to the correct path.
+    return "del"

pdd/get_extension.py

@@ -1,7 +1,7 @@
 """Module to retrieve file extensions for programming languages."""
 
-import os
 import pandas as pd
+from pdd.path_resolution import get_default_resolver
 
 def get_extension(language: str) -> str:
     """
@@ -18,13 +18,12 @@ def get_extension(language: str) -> str:
         ValueError: If the PDD_PATH environment variable is not set.
         FileNotFoundError: If the language_format.csv file is not found.
     """
-    # Step 1: Load the environment variable PDD_PATH
-    pdd_path = os.getenv('PDD_PATH')
-    if not pdd_path:
-        raise ValueError("Environment variable PDD_PATH is not set.")
-    
-    # Construct the full path to the CSV file
-    csv_file_path = os.path.join(pdd_path, 'data', 'language_format.csv')
+    # Step 1: Resolve CSV path from PDD_PATH
+    resolver = get_default_resolver()
+    try:
+        csv_file_path = resolver.resolve_data_file("data/language_format.csv")
+    except ValueError as exc:
+        raise ValueError("Environment variable PDD_PATH is not set.") from exc
     
     # Step 2: Lower case the language string
     language_lower = language.lower()
@@ -63,4 +62,4 @@ def get_extension(language: str) -> str:
 #   ```bash
 #   pip install pandas
 #   ```
-# - Ensure that the CSV file is structured correctly and located at the specified path.
+# - Ensure that the CSV file is structured correctly and located at the specified path.