Automated Python Project Generator with Self-Healing Tests

main.py

import subprocess
import time
import re
import shutil
from pathlib import Path

import openai
from slugify import slugify

api_key = "sk-proj-..."  # Replace with your OpenAI API key

client = openai.OpenAI(api_key=api_key)

MODEL = "gpt-4-turbo"


def ask_task() -> str:
    return input("💬 What do you want to build? ").strip()


def run(cmd: list[str], cwd: Path | None = None) -> None:
    subprocess.run(cmd, cwd=cwd, check=True)


def run_capture(cmd: list[str], cwd: Path | None = None) -> tuple[bool, str]:
    process = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
    return process.returncode == 0, process.stdout + process.stderr


def extract_code(text: str) -> str:
    match = re.search(r"```(?:python)?(.*?)```", text, re.DOTALL)
    return match.group(1).strip() if match else text.strip()


def call_openai(messages: list[dict[str, str]]) -> str:
    response = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        temperature=0,
    )
    return extract_code(response.choices[0].message.content)


def infer_project_metadata(task: str) -> tuple[str, str]:
    messages = [
        {"role": "system", "content": "You are a Python project assistant."},
        {"role": "user", "content": f"For the task: '{task}', suggest:\n"
                                    "- a readable project name (dashes allowed)\n"
                                    "- a valid Python module name (snake_case)\n"
                                    "Return in the format: project-name module_name"}
    ]
    result = call_openai(messages)

    lines = result.strip().splitlines()
    values = []

    for line in lines:
        if ':' in line:
            values.append(line.split(":", 1)[1].strip())
        elif line.strip():
            values.extend(line.strip().split())

    if len(values) >= 2:
        project = slugify(values[0], separator='-')
        module = slugify(values[1], separator='_')
        return project, module

    raise ValueError(f"Unexpected response from OpenAI: {result}")


def infer_dependencies(task: str) -> list[str]:
    messages = [
        {"role": "system", "content": "You are a Python build assistant."},
        {"role": "user", "content": f"What Python packages (PyPI) are required to implement this task: '{task}'? Return only the list separated by spaces."}
    ]
    deps = call_openai(messages)
    return deps.split()


def generate_files(task: str, project_dir: Path, module_name: str) -> None:
    impl_filename = module_name  # enforce .py filename == module name

    src_path = project_dir / "src" / module_name
    tests_path = project_dir / "tests"
    src_path.mkdir(parents=True, exist_ok=True)
    tests_path.mkdir(parents=True, exist_ok=True)

    # empty __init__.py
    (src_path / "__init__.py").write_text("")

    # generate test file
    test_path = tests_path / f"test_{impl_filename}.py"
    messages_test = [
        {"role": "system", "content": "You write a test file using pytest."},
        {"role": "user", "content": f"Write a pytest test file for the implementation in src/{module_name}/{impl_filename}.py.\n"
                                    f"Only write tests. Do not include the implementation itself."}
    ]
    test_code = call_openai(messages_test)
    test_path.write_text(test_code)

    # generate implementation
    impl_path = src_path / f"{impl_filename}.py"
    messages_code = [
        {"role": "system", "content": "You write a Python implementation module."},
        {"role": "user", "content": f"Write the implementation for the following task:\n\n{task}\n\n"
                                    f"The file is: src/{module_name}/{impl_filename}.py\n"
                                    f"Only return the implementation. Do not include tests."}
    ]
    impl_code = call_openai(messages_code)
    impl_path.write_text(impl_code)


def fix_code(current_code: str, feedback: str) -> str:
    messages = [
        {"role": "system", "content": "You fix the code to make the tests pass."},
        {"role": "user", "content": f"Here is the current code:\n\n```python\n{current_code}\n```"},
        {"role": "user", "content": f"The following pytest output shows errors:\n\n{feedback}"}
    ]
    return call_openai(messages)


def main() -> None:
    task = ask_task()
    project_name, module_name = infer_project_metadata(task)
    project_dir = Path(project_name)

    if project_dir.exists():
        shutil.rmtree(project_dir)

    print(f"🚀 Creating project: {project_name}")
    run(["uv", "init", "--lib", project_name])

    print("📦 Detecting dependencies...")
    deps = ["pytest"] + infer_dependencies(task)
    print(f"📦 Dependencies: {deps}")
    run(["uv", "add"] + deps, cwd=project_dir)

    print("📝 Generating code and tests...")
    generate_files(task, project_dir, module_name)

    code_path = project_dir / "src" / module_name / f"{module_name}.py"

    iteration = 0
    while True:
        print(f"\n🔁 Iteration {iteration}")
        success, output = run_capture(["uv", "run", "python", "-m", "pytest", "-q", "--tb=short"], cwd=project_dir)

        if success:
            print("✅ All tests passed.")
            break

        print("❌ Tests failed. Attempting fix...")
        current_code = code_path.read_text()
        new_code = fix_code(current_code, output)
        code_path.write_text(new_code)

        time.sleep(1)
        iteration += 1


if __name__ == "__main__":
    main()

$ uv run python main_openai.py 
💬 What do you want to build? a simple linear regression model
🚀 Creating project: simple-linear-regression
Adding `simple-linear-regression` as member of workspace `/home/goabonga/autocode`
Initialized project `simple-linear-regression` at `/home/goabonga/autocode/simple-linear-regression`
📦 Detecting dependencies...
📦 Dependencies: ['pytest', 'numpy', 'scipy', 'scikit-learn']
Resolved 36 packages in 89ms
      Built simple-linear-regression @ file:///home/goabonga/autocode/simple-linear-regression
Prepared 1 package in 295ms
Installed 1 package in 0.58ms
 + simple-linear-regression==0.1.0 (from file:///home/goabonga/autocode/simple-linear-regression)
📝 Generating code and tests...

🔁 Iteration 0
❌ Tests failed. Attempting fix...

🔁 Iteration 1
✅ All tests passed.

$ tree simple-linear-regression/
simple-linear-regression/
├── pyproject.toml
├── README.md
├── src
│   └── simple_linear_regression
│       ├── __init__.py
│       ├── py.typed
│       └── simple_linear_regression.py
└── tests
    └── test_simple_linear_regression.py

test_simple_linear_regression.py :

import pytest
import numpy as np
from src.simple_linear_regression.simple_linear_regression import SimpleLinearRegression

def test_initialization():
    """Test the initialization of the SimpleLinearRegression model."""
    model = SimpleLinearRegression()
    assert model.slope_ is None, "Initial slope should be None"
    assert model.intercept_ is None, "Initial intercept should be None"

def test_fit():
    """Test fitting the model with simple linear data."""
    # Create a simple linear relationship y = 2x + 1
    X = np.array([1, 2, 3, 4, 5])
    y = 2 * X + 1
    model = SimpleLinearRegression()
    model.fit(X, y)
    
    assert model.slope_ == 2, "Slope should be correctly calculated"
    assert model.intercept_ == 1, "Intercept should be correctly calculated"

def test_predict():
    """Test predictions from the fitted model."""
    X = np.array([1, 2, 3, 4, 5])
    y = 2 * X + 1
    model = SimpleLinearRegression()
    model.fit(X, y)
    
    predictions = model.predict(np.array([6, 7]))
    expected = np.array([13, 15])  # Since y = 2x + 1
    
    np.testing.assert_array_equal(predictions, expected, "Predictions should match expected values")

def test_fit_invalid_input():
    """Test fitting the model with invalid inputs."""
    model = SimpleLinearRegression()
    with pytest.raises(ValueError):
        model.fit(np.array([1, 2, 3]), np.array([1, 2]))  # Mismatched lengths

def test_predict_without_fit():
    """Test prediction before fitting the model."""
    model = SimpleLinearRegression()
    with pytest.raises(ValueError):
        model.predict(np.array([1, 2]))

# Optionally, you can add more complex scenarios or edge cases

simple_linear_regression.py :

# src/simple_linear_regression/simple_linear_regression.py

import numpy as np

class SimpleLinearRegression:
    def __init__(self):
        self.slope_ = None
        self.intercept_ = None

    def fit(self, X, y):
        """
        Fit the simple linear regression model to the training data.
        
        Parameters:
            X (array-like): 1D array of feature values.
            y (array-like): 1D array of target values.
        """
        X = np.array(X)
        y = np.array(y)
        
        # Calculate the mean of X and y
        x_mean = np.mean(X)
        y_mean = np.mean(y)
        
        # Calculate the terms needed for the numerator and denominator of beta
        numerator = np.sum((X - x_mean) * (y - y_mean))
        denominator = np.sum((X - x_mean) ** 2)
        
        # Calculate slope (beta)
        self.slope_ = numerator / denominator
        
        # Calculate intercept (alpha)
        self.intercept_ = y_mean - self.slope_ * x_mean

    def predict(self, X):
        """
        Predict the target values using the linear model.
        
        Parameters:
            X (array-like): 1D array of feature values.
        
        Returns:
            array: Predicted target values.
        """
        if self.slope_ is None or self.intercept_ is None:
            raise ValueError("Model has not been fitted yet.")
        
        X = np.array(X)
        return self.slope_ * X + self.intercept_

    def get_slope(self):
        """
        Returns the slope of the linear model.
        
        Returns:
            float: The slope of the model.
        """
        return self.slope_

    def get_intercept(self):
        """
        Returns the intercept of the linear model.
        
        Returns:
            float: The intercept of the model.
        """
        return self.intercept_