{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f705f4be70e9"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j1cu3boIga_X"
      },
      "source": [
        "# Reduce Tech Debt with Gemini 3 Pro\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/refactor_legacy_code.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fuse-cases%2Fcode%2Frefactor_legacy_code.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/use-cases/code/refactor_legacy_code.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/refactor_legacy_code.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/refactor_legacy_code.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://storage.googleapis.com/github-repo/generative-ai/logos/GitHub_Invertocat_Dark.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<p>\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/refactor_legacy_code.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/refactor_legacy_code.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/refactor_legacy_code.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/refactor_legacy_code.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/refactor_legacy_code.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>\n",
        "</p>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "83f287773529"
      },
      "source": [
        "| Authors |\n",
        "| --- |\n",
        "| [Dennis Kashkin](https://github.com/dkashkin) |\n",
        "| [Holt Skinner](https://github.com/holtskinner) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "308adc2720c1"
      },
      "source": [
        "## Overview: Escaping \"Legacy Land\"\n",
        "\n",
        "We all know the code. The ten-year-old library that nobody quite understands anymore. You don't want to be the one who breaks it, but you can't ignore it forever. The struggle is so relatable that it inspired [Forrest, Amanda, and Holt to sing about it at Next '23](https://goo.gle/legacy-land).\n",
        "\n",
        "But while the fear is real, so is the need for modernization. Security, performance, and scalability demand that we refactor.\n",
        "\n",
        "**About this Notebook**\n",
        "We are going to look at a new approach to this old problem. utilizing **Gemini 3 Pro on Vertex AI**. We will look at how to leverage the model not just for faster coding, but for analyzing dependencies and executing safer, more reliable code modernization."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VbTqt86rrbWQ"
      },
      "source": [
        "## Prerequisites\n",
        "\n",
        "Before starting, ensure you have the necessary dependencies and environment set up.\n",
        "\n",
        "*(Note: In a Jupyter environment, you may need to restart the kernel after installing new packages.)*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OzilzKKFrd2m"
      },
      "outputs": [],
      "source": [
        "%pip install -q --upgrade google-genai pyparseit"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XRGl-_glga_Z"
      },
      "source": [
        "### Authenticate your notebook environment\n",
        "\n",
        "If you are running this notebook in **Google Colab**, run the cell below to authenticate your account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rK-IZVyzga_Z"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b6d46b3a5abf"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "05dab3326853"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from google import genai\n",
        "\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "LOCATION = \"global\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "# fmt: on\n",
        "\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "if not LOCATION:\n",
        "    LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\")\n",
        "\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "935c5cca7fb3"
      },
      "source": [
        "### Load model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6178371d0b7a"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"gemini-3-pro-preview\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bc778f7deef4"
      },
      "source": [
        "### Import libraries and create helper functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "778300e059d9"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import logging\n",
        "import re\n",
        "\n",
        "import markdown\n",
        "from IPython.display import HTML, display\n",
        "from google.genai import types\n",
        "from pyparseit import parse_markdown_string\n",
        "\n",
        "logging.disable(level=logging.WARNING)\n",
        "\n",
        "\n",
        "def call_llm(prompt: str) -> str:\n",
        "    \"\"\"Simple wrapper for model inference.\"\"\"\n",
        "    print(f\"Calling {MODEL_ID} with {len(prompt)} chars...\")\n",
        "    response = client.models.generate_content(\n",
        "        model=MODEL_ID,\n",
        "        contents=prompt,\n",
        "        config=types.GenerateContentConfig(\n",
        "            temperature=1.0,\n",
        "            thinking_config=types.ThinkingConfig(thinking_budget=16384),\n",
        "        ),\n",
        "    )\n",
        "    # Basic cleanup of markdown code blocks if present\n",
        "    snippets = parse_markdown_string(response.text, language=\"python\")\n",
        "    return snippets[0].content\n",
        "\n",
        "\n",
        "def generate_and_save(template_path: str, output_path: str, replacements: dict) -> str:\n",
        "    \"\"\"Helper to load a prompt template, inject variables, call LLM, and save the result.\"\"\"\n",
        "    with open(template_path) as f:\n",
        "        prompt = f.read()\n",
        "\n",
        "    for key, value in replacements.items():\n",
        "        prompt = prompt.replace(f\"{{{key}}}\", value)\n",
        "\n",
        "    print(f\"Generating from {template_path}...\")\n",
        "    content = call_llm(prompt)\n",
        "\n",
        "    with open(output_path, \"w\") as f:\n",
        "        f.write(content)\n",
        "    return content\n",
        "\n",
        "\n",
        "def render_markdown_with_mermaid(markdown_string: str) -> None:\n",
        "    \"\"\"Renders a string containing markdown and mermaid diagrams inline.\"\"\"\n",
        "    # Placeholder and storage logic is identical to the cell magic\n",
        "    mermaid_placeholder = \"---MERMAID_PLACEHOLDER_{}---\"\n",
        "    mermaid_blocks = re.findall(r\"```mermaid\\n(.*?)\\n```\", markdown_string, re.DOTALL)\n",
        "\n",
        "    temp_md = markdown_string\n",
        "    for i, block in enumerate(mermaid_blocks):\n",
        "        temp_md = temp_md.replace(\n",
        "            f\"```mermaid\\n{block}\\n```\", mermaid_placeholder.format(i), 1\n",
        "        )\n",
        "\n",
        "    html_output = markdown.markdown(\n",
        "        temp_md, extensions=[\"extra\", \"fenced_code\", \"tables\"]\n",
        "    )\n",
        "\n",
        "    for i, block in enumerate(mermaid_blocks):\n",
        "        encoded_mermaid = base64.b64encode(block.encode(\"utf-8\")).decode(\"utf-8\")\n",
        "        mermaid_html = f\"\"\"\n",
        "        <div class='mermaid' style='margin-top: 1em; margin-bottom: 1em;'><iframe\n",
        "            src=\"https://mermaid.ink/svg/{encoded_mermaid}\"\n",
        "            width=\"100%\" height=\"450\" frameborder=\"0\" scrolling=\"no\"\n",
        "            style=\"max-width: 100%; display: block; margin: auto;\"\n",
        "        ></iframe></div>\"\"\"\n",
        "        html_output = html_output.replace(mermaid_placeholder.format(i), mermaid_html)\n",
        "\n",
        "    display(HTML(html_output))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2At_vpZvga_b"
      },
      "source": [
        "### Create necessary directories."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gQhCIhZCYxs2"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "DIRS = [\n",
        "    \"original_code\",\n",
        "    \"original_tests\",\n",
        "    \"output/tests\",\n",
        "    \"output/specs\",\n",
        "    \"output/design_docs\",\n",
        "    \"output/code\",\n",
        "    \"prompts\",\n",
        "]\n",
        "for d in DIRS:\n",
        "    os.makedirs(d, exist_ok=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uTlfZ5Grga_b"
      },
      "source": [
        "#### Load Prompts\n",
        "\n",
        "Load the specialized prompts for each stage of the pipeline."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dOoN54WMga_b"
      },
      "outputs": [],
      "source": [
        "%%writefile prompts/create_tests.md\n",
        "<role>\n",
        "You are a world-class software engineer specialized in writing comprehensive unit tests for specific Python modules.\n",
        "</role>\n",
        "\n",
        "<task>\n",
        "Your goal is to write a robust set of \"Passing Tests\" for the Python module provided in the `source_code` tag below.\n",
        "\"Passing Tests\" are unit tests that cover existing, valid code paths and are expected to pass.\n",
        "Do not write any \"Failing Tests\" that are expected to result in errors based on the current version of source code. For example, if you see insufficient validation of input parameters or lack of exception handling, do not write any tests that target such known gaps.\n",
        "</task>\n",
        "\n",
        "<instructions>\n",
        "You must follow these instructions:\n",
        "- You must use Python's standard framework \"unittest\"\n",
        "- You must write a complete Python module including all required import statements\n",
        "- Wrap all unit tests into one test class inherited from unittest.TestCase\n",
        "- The source code of the tested module must be imported using the following exact code:\n",
        "\n",
        "```python\n",
        "import importlib\n",
        "import os\n",
        "import sys\n",
        "\n",
        "sys.path.insert(0, os.getcwd())\n",
        "PACKAGE_PREFIX = os.environ.get('PACKAGE_PREFIX', '{code_dir_path}')\n",
        "module = importlib.import_module(f'{PACKAGE_PREFIX}.{module_name}')\n",
        "```\n",
        "\n",
        "- Include the following code at the very end of the Python module with the unit test class:\n",
        "\n",
        "```python\n",
        "if __name__ == '__main__':\n",
        "    unittest.main()\n",
        "```\n",
        "\n",
        "</instructions>\n",
        "\n",
        "<response_format>\n",
        "Format your response as a single Markdown code block with the full source code of the Python module containing all unit tests.\n",
        "</response_format>\n",
        "\n",
        "<best_practices>\n",
        "# A Comprehensive Guide to Unit Testing Python Modules\n",
        "\n",
        "Writing effective unit tests is a critical skill in software engineering. When testing a Python module with multiple classes and functions, the scope expands from testing individual functions in isolation to verifying the collaborative behavior of those methods and managing the state of class instances. A key metric for a thorough test suite is code coverage, which measures the percentage of your code executed by your tests. While aiming for 100% coverage is an excellent goal that builds confidence and reduces bugs, remember that it is a tool, not the ultimate goal. A meaningful test that verifies the class's public behavior is always more valuable than a test that simply executes a line of code.\n",
        "\n",
        "## 1. Understand the Role and Public Interface of every class and function\n",
        "\n",
        "Before writing a single test, you must thoroughly understand the classes and functions you intend to test. This involves looking at the bigger picture of how each class and its methods are intended to be used.\n",
        "\n",
        "*   **Analyze the Responsibilities**: Read the code, docstrings, and any related specifications to understand the purpose, the public methods (the intended interface), and the interactions with other components.\n",
        "*   **Focus on Public Behavior**: Prioritize testing the public methods of classes, as these define how other parts of your code will interact with it. While it's possible to test private methods in Python, it's often a sign that your class may have too many responsibilities and could be a candidate for refactoring. Testing through the public interface helps ensure your tests aren't tightly coupled to implementation details, making them less brittle to future code changes.\n",
        "*   **Identify Method Interactions and State Changes**: Map out how the methods within each class interact with each other. Does one method's execution affect the state of the object, which in turn influences the outcome of another method? These interactions are crucial to test.\n",
        "*   **Map Out All Code Paths**: For each public method, and any critical private methods that contain complex logic, you must identify and test all execution paths created by conditional logic.\n",
        "    *   **Conditionals**: Every `if`, `elif`, and `else` statement creates a branch that needs to be tested.\n",
        "    *   **Loops**: Test scenarios where a loop runs zero, one, or multiple times. Also, test any `break` or `continue` statements that alter the loop's flow.\n",
        "    *   **Exception Handling**: Ensure your tests trigger each `try` block's success path and every corresponding `except` block's failure path.\n",
        "    *   **Early Returns**: Functions with guard clauses or multiple return statements have exits that must all be tested.\n",
        "\n",
        "## 2. Design Test Cases Systematically\n",
        "\n",
        "A structured approach to test case design ensures that you cover all expected behaviors, boundaries, and failure scenarios.\n",
        "\n",
        "*   **Equivalence Partitioning**: Divide possible inputs for each method into groups that should be handled similarly. For example, a method processing a user's age might have partitions for valid ages (e.g., 18-65), underage (e.g., 0-17), and invalid inputs (e.g., negative numbers or non-integer types). Test one representative from each group.\n",
        "*   **Boundary Value Analysis**: Pay special attention to the edges of your input ranges. For an age range of 18-65, you should test the values 17, 18, 65, and 66.\n",
        "*   **Test the \"Happy Path\"**: Start with at least one test for the normal, expected use case of each public method with valid inputs. For a class, this may involve a sequence of method calls that represent a typical interaction.\n",
        "*   **Cover Edge Cases and Invalid Inputs**: This is where many bugs hide. Your tests should cover:\n",
        "    *   **Empty or Null Inputs**: `None`, empty strings (`\"\"`), and empty collections (`[]`, `{}`).\n",
        "    *   **Numeric Boundaries**: `0`, `-1`, and very large or small numbers.\n",
        "    *   **Invalid Types**: Pass a string where an integer is expected to ensure proper error handling.\n",
        "*   **Test State Transitions**: For stateful classes, design tests that specifically check the transition of the object's state after a method is called. For instance, if a `user.deactivate()` method is called, a subsequent call to `user.is_active()` should return `False`.\n",
        "\n",
        "## 3. Leverage Mocking\n",
        "\n",
        "Classes and functions often depend on external systems like databases, APIs, or the file system. To test your logic in isolation, you must use **mocking**. Mocking involves replacing these external dependencies with controlled objects that simulate their behavior.\n",
        "\n",
        "*   **Why Mock?**: Mocking makes your tests faster, more reliable, and independent of external systems. It allows you to simulate success, failure (e.g., an API timeout), and edge case responses from dependencies without making real network or disk operations.\n",
        "*   **How to Mock**: Python's `unittest.mock` library provides tools like `@patch` to replace objects within a test's scope. When mocking, it's a good practice to use `autospec=True` to ensure your mock has the same interface as the object it's replacing.\n",
        "*   **Dependency Injection**: A powerful technique to make your functions and classes more testable is dependency injection. Instead of a class creating its own dependencies internally, you provide them as arguments to the constructor. This makes it easy to substitute a real dependency with a mock during testing.\n",
        "\n",
        "**Example: Mocking an API Call in a Class Method**\n",
        "\n",
        "```python\n",
        "from unittest.mock import patch, Mock\n",
        "\n",
        "# Assuming 'user_service.get_user_data' calls 'requests.get'\n",
        "@patch('my_app.user_service.requests.get')\n",
        "def test_get_user_data_success(mock_get):\n",
        "    # Configure the mock to return a successful response\n",
        "    mock_response = Mock()\n",
        "    mock_response.json.return_value = {'id': 1, 'name': 'Jane Doe'}\n",
        "    mock_get.return_value = mock_response\n",
        "\n",
        "    # Instantiate the service and call the method\n",
        "    service = UserService()\n",
        "    user_data = service.get_user_data(1)\n",
        "\n",
        "    # Assert the result\n",
        "    assert user_data['name'] == 'Jane Doe'\n",
        "```\n",
        "\n",
        "## 4. Recommended Workflow\n",
        "\n",
        "1.  **Write Initial Tests**: Begin by writing tests for the primary functionality, focusing on the public methods and their expected outcomes (the \"happy path\").\n",
        "2.  **Manage Test State with `setUp` and `tearDown`**: Use the `setUp` method in your test class to create a clean instance before each test method runs. This ensures that your tests are independent and don't influence each other. The `tearDown` method can be used to clean up any resources after each test.\n",
        "3.  **Identify Gaps**: Review the code to find uncovered lines, branches, and method interactions.\n",
        "4.  **Write Targeted Tests**: For each missed scenario, determine the specific inputs, method call sequences, or mock behaviors needed to execute that code path. Write a new test case targeting that exact situation.\n",
        "5.  **Refactor for Clarity**: As you add tests, keep them clean and readable. Use descriptive names (e.g., `test_user_deactivation_sets_is_active_to_false`) and consider using parameterized tests to avoid duplicating test logic when testing similar inputs.\n",
        "\n",
        "## 5. Advanced Techniques and Best Practices\n",
        "\n",
        "1.  **Parameterize Tests**: To test a method with multiple inputs using the same logic, use a loop with `subTest` or a parameterized testing library like `pytest`. This keeps your test suite DRY (Don't Repeat Yourself).\n",
        "2.  **Verify State and Side Effects**: A test should not only check a method's return value but also verify any intended side effects, such as changes to the object's attributes or calls to a mocked dependency.\n",
        "3.  **Test for Exceptions**: Use `assertRaises` (as a context manager or method) to confirm that a method correctly raises an exception with invalid input.\n",
        "4.  **Keep Tests Independent**: Each test case should be able to run on its own and in any order. Avoid creating dependencies between tests where one test relies on the state left by a previous one.\n",
        "5.  **Use Descriptive Naming**: Test methods should have long, descriptive names that clearly state what they are testing. This makes test results easier to understand and serves as a form of documentation.\n",
        "</best_practices>\n",
        "\n",
        "<source_code>\n",
        "{source_code}\n",
        "</source_code>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I7r4Di-Xga_c"
      },
      "outputs": [],
      "source": [
        "%%writefile prompts/create_spec.md\n",
        "\n",
        "<role>\n",
        "You are a world-class software architect specialized in designing comprehensive specifications intended for junior developers.\n",
        "</role>\n",
        "\n",
        "<task>\n",
        "Review the source code of the Python module provided in the `source_code` tag below and write a comprehensive Software Requirements Specification for this module by strictly following the process defined in the \"process\" tag below.\n",
        "</task>\n",
        "\n",
        "<constraints>\n",
        "1. The specification must be written in plain English language, without any Python code snippets.\n",
        "2. You can include diagrams in Mermaid format.\n",
        "3. You must follow the Specification Design Process provided in the `process` tag below.\n",
        "</constraints>\n",
        "\n",
        "<process>\n",
        "## Guiding Principles\n",
        "\n",
        "1.  **Strict Contract**: The resulting SRS is a non-negotiable contract. Deviations in the final implementation, even for technical improvement, are considered defects.\n",
        "2.  **Independent Resolution**: You MUST resolve all ambiguities autonomously based on currently available information (code, comments, tests). Do not rely on external clarification. Make reasonable engineering assumptions based on the simplest non-crashing behavior or standard patterns to ensure deterministic outcomes.\n",
        "3.  **Absolute Completeness**: No placeholders (e.g., \"etc.\", \"such as\", \"e.g.\"). Every list, enum, dictionary key set, attribute, method, and constant MUST be exhaustively enumerated.\n",
        "4.  **Precision and Clarity**: Avoid vague terms. Define any domain-specific terms used. Specifications MUST be unambiguous and specific.\n",
        "\n",
        "## Phase 1: Scope Definition and Strategic Planning\n",
        "\n",
        "Establish the boundaries of the documentation effort, ensuring no part of the exposed surface area is omitted.\n",
        "\n",
        "1.  **Define Core Objective**: Produce a complete SRS for all components within the target Python module.\n",
        "2.  **Identify the Complete Module Namespace**: The scope is strictly defined by the *entire* module namespace. This includes ALL publicly accessible components, protected components, and any private components accessed by the test suite:\n",
        "    *   Classes and their full inheritance hierarchies.\n",
        "    *   Module-level functions and helper utilities.\n",
        "    *   Exceptions, constants, type variables, and protocol definitions.\n",
        "    *   Any internal component (even those prefixed with `_`) if it is accessed by the existing test suite, treating such access as defining a \"hidden\" public contract.\n",
        "    *   Required module-level imports, including from `__future__`.\n",
        "3.  **Conditional Features**: Identify features or classes available only under certain conditions (e.g., OS, Python version, presence of other modules). The SRS MUST specify exactly how their absence is reflected in the module namespace. These components SHALL NOT be defined in the namespace if the condition is not met.\n",
        "    *   *Example:* \"`KqueueSelector` class SHALL be present in the module namespace if and only if `select.kqueue` is available.\"\n",
        "\n",
        "## Phase 2: Static Analysis and Architectural Deconstruction\n",
        "\n",
        "Examine the source code without execution to map its structure, contracts, and dependencies with absolute precision.\n",
        "\n",
        "1.  **Namespace and Hierarchy Cataloging**:\n",
        "    *   Create a definitive list of every name exported by the module.\n",
        "    *   Formalize all class hierarchies, explicitly documenting required naming conventions and exact inheritance trees.\n",
        "    *   List all required module-level constants with their types and expected values.\n",
        "        *   *Example:* `DEFAULT_TIMEOUT: final[int] = 10`\n",
        "    *   List all required imports, including from `__future__`. Verify the existence and exact names of any functions or attributes assumed to be imported from other modules.\n",
        "        *   *Example:* \"SHALL import `annotations` from `__future__`.\"\n",
        "    *   *Constraint Check*: Ensure NO \"etc.\" or \"e.g.\" is used. All items MUST be listed.\n",
        "\n",
        "2.  **Comprehensive API Extraction**:\n",
        "    *   **Classes**:\n",
        "        *   List all attributes (public, protected, private): Specify name, type, access level (public, protected, private), and if it's a property or a plain instance/class attribute.\n",
        "            *   *Example:* `_is_open: bool` - Instance Attribute, Private, Read/Write.\n",
        "            *   *Example:* `@property max_size(self) -> int` - Property, Public, Read-Only.\n",
        "        *   List all methods (public, protected, private): Provide the full signature, including parameter types, return types, and decorators (`@staticmethod`, `@classmethod`). Pay close attention to dunder methods, ensuring correct parameters (e.g., `self`).\n",
        "            *   *Example:* `def _process_data(self, data: bytes, config: Dict[str, Any]) -> bool:`\n",
        "            *   *Example:* `def __len__(self) -> int:`\n",
        "    *   **Module-Level Callables (Functions)**: Extract the *exact* signature for every callable, including type hints for all parameters and return values. Explicit documentation for `*args` and `**kwargs` handling is mandatory.\n",
        "    *   **Data Structures**: Deconstruct all Data Transfer Objects (DTOs), enums, and configuration dictionaries. Every allowed key or member MUST be enumerated. Use type hints to define the structure of complex types.\n",
        "        *   *Example:* Config dict structure: `{\"timeout\": int, \"retries\": int, \"protocol\": Literal[\"tcp\", \"udp\"]}`\n",
        "        *   Specify mutability where important (e.g., `set` vs. `frozenset`, `List` vs. `Tuple`).\n",
        "\n",
        "## Phase 3: Behavioral Analysis and Protocol Reverse-Engineering\n",
        "\n",
        "Understand dynamic behavior, state transitions, and implicit Python protocols through rigorous examination of code and tests.\n",
        "\n",
        "1.  **Test Suite Deconstruction as Source of Truth**:\n",
        "    *   Analyze every test case to extract hard input-output scenarios and edge case handling.\n",
        "    *   Treat mocked interactions in tests as rigid requirements for external dependencies.\n",
        "    *   Identify \"hidden\" public contracts: If a test accesses an attribute named `_internal_cache`, that attribute MUST be documented as part of the specification as if it were public for testing purposes.\n",
        "\n",
        "2.  **Implicit Protocol Identification**:\n",
        "    *   Identify all implemented \"dunder\" methods (e.g., `__add__`, `__iter__`, `__call__`, `__getstate__`, `__str__`, `__repr__`).\n",
        "    *   Document these not just as methods, but as formal adherence to Python protocols (e.g., \"Implements the `Iterable` protocol via `__iter__`\").\n",
        "\n",
        "3.  **Logic, Side Effect, and Algorithm Mapping**:\n",
        "    *   Use clear, sequential steps (e.g., numbered lists) and imperative language (\"SHALL\", \"MUST\") to describe method/function behavior.\n",
        "    *   Specify the order of operations, especially when interactions or side effects are involved.\n",
        "        *   *Example:* \"1. SHALL validate input `config`. 2. SHALL acquire lock. 3. SHALL call `_internal_method()` *before* releasing lock.\"\n",
        "    *   Enumerate all known edge cases and specify exactly how they MUST be handled.\n",
        "        *   *Example:* \"If input is None, SHALL raise `TypeError`. If input list is empty, SHALL return an empty `Dict`.\"\n",
        "    *   Provide concrete examples of inputs and expected outputs, especially for parsing, formatting, or complex transformations.\n",
        "    *   Document ALL observable side effects exactly (e.g., specific log message templates, exact file system paths utilized, mutations to input objects if applicable).\n",
        "    *   **String Formats**: For methods returning strings (`__str__`, `__repr__`, custom formatting), provide the exact format string or a template.\n",
        "        *   *Example:* `__repr__` SHALL return `f\"MyClass(id={self.id!r})\"`\n",
        "    *   **Regular Expressions**: Include any required regular expressions in their entirety, with flags. Explain capture groups.\n",
        "        *   *Example:* `VERSION_REGEX = re.compile(r\"V(?P<major>\\d+)\", re.IGNORECASE)`\n",
        "    *   **Algorithms/Formulas**: If a specific algorithm is required, provide step-by-step pseudo-code. Explicitly state any mathematical formulas. Detail data transformations like sorting, stripping, encoding, or type conversions.\n",
        "        *   *Example:* \"Output from command SHALL be decoded as UTF-8, stripped of leading/trailing whitespace, and split by newline characters.\"\n",
        "\n",
        "## Phase 4: Specification Synthesis\n",
        "\n",
        "Consolidate findings into a formal SRS document. This phase requires making definitive decisions to resolve all ambiguities without external input.\n",
        "\n",
        "1.  **Ambiguity Resolution and Assumption Locking**:\n",
        "    *   Explicitly document every ambiguity encountered during analysis.\n",
        "    *   Resolve each ambiguity by making and documenting a reasonable engineering assumption. Base assumptions on promoting robustness, simplicity, deterministic behavior, or adherence to common Python idioms.\n",
        "    *   *Example Resolution*: \"The code handles `None` for `user_id` inconsistently. SRS mandates `user_id` MUST NOT be None and SHALL raise `ValueError` if None is provided.\"\n",
        "    *   Avoid vague terms. Define any domain-specific terms used.\n",
        "\n",
        "2.  **SRS Structural Requirements**:\n",
        "    The SRS MUST include standard sections (Introduction, Overview, NFRs) AND specifically mandated sections:\n",
        "    *   **Module Namespace Reference**: A dedicated section defining every exposed name (classes, functions, exceptions, constants), including conditional components as defined in Phase 1.\n",
        "    *   **Assumption Log**: A record of all independent decisions made to resolve established ambiguities, with justifications.\n",
        "\n",
        "3.  **Detailed Component Specification Rules**:\n",
        "    Every component MUST be documented using the following strict formats:\n",
        "\n",
        "    *   **For Classes**:\n",
        "        *   **Inheritance**: The exact parent class(es).\n",
        "        *   **Attributes**: A complete table of ALL public, protected, and private attributes (Name, Type, Access Level, Description, Property/Instance/Class).\n",
        "            *   *Example:* `is_initialized: bool` | Instance Attribute | Private | Tracks initialization state.\n",
        "        *   **Methods**: Detailed specification for each method including dunder methods (see Callables below).\n",
        "        *   **Protocol Support**: Explicit list of supported Python protocols.\n",
        "\n",
        "    *   **For Callables (Methods/Functions)**:\n",
        "        *   **Full Signature**: MUST include all parameters with type hints and default values. Decorators MUST be listed.\n",
        "            *   *Example:* `@classmethod def from_config(cls, config: Dict[str, Any]) -> 'MyClass':`\n",
        "        *   **`*args`/`**kwargs` definition**: Explicit explanation of how variadic arguments are processed, including expected types and structure if applicable.\n",
        "        *   **Preconditions**: What MUST be true before execution.\n",
        "            *   *Example:* \"Parameter `config` MUST not be None.\"\n",
        "        *   **Behavior**: Detailed step-by-step description using imperative language, including order of operations, side effects, algorithms, and data transformations as per Phase 3.3.\n",
        "        *   **Postconditions**: What is guaranteed to be true after successful execution.\n",
        "        *   **Return Value**: Exact type and structure. For complex types, provide a breakdown.\n",
        "            *   *Example:* Returns `Dict[str, List[Tuple[int, str]]]` where each tuple is `(index, value)`.\n",
        "        *   **Error Handling**: Exact exception types and their *literal* error message strings (use f-string like placeholders for dynamic parts) for every error condition.\n",
        "            *   *Example:* SHALL raise `ValueError` with message `f\"Invalid ID: {id}\"` if ID is not positive.\n",
        "            *   *Example:* SHALL raise `TypeError` with message \"Configuration must be a dict.\"\n",
        "\n",
        "    *   **For Collections/Enums/Constants**:\n",
        "        *   **Complete Enumeration**: ABSOLUTELY NO use of \"e.g.\" or \"etc.\". Every member, key, or value MUST be listed.\n",
        "        *   **Type and Value**: For constants, specify type and exact value.\n",
        "\n",
        "## Phase 5: Verification and Validation\n",
        "\n",
        "Ensure the SRS is accurate, complete, and actionable based *solely* on the generated document.\n",
        "\n",
        "1.  **Canonical Test Suite Definition**:\n",
        "    *   Translate identified behavioral scenarios, edge cases, and input/output examples from Phase 3 into a required acceptance test suite definition within the SRS.\n",
        "    *   Ensure these tests cover the \"hidden contracts\" (e.g., `_private` attributes accessed by legacy tests) as specified.\n",
        "\n",
        "2.  **Completeness Review**:\n",
        "    *   Verify that NO placeholders (\"etc.\", \"e.g.\") exist anywhere in the document.\n",
        "    *   Validate that every name in the module namespace is documented.\n",
        "    *   Confirm that every assumption made in Phase 4 is clearly documented in the Assumption Log and reflected as a firm requirement in the component specifications.\n",
        "    *   Ensure all types, string formats, error messages, and behaviors are specified with an absolute level of detail.\n",
        "</process>\n",
        "\n",
        "<source_code>\n",
        "{source_code}\n",
        "</source_code>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EHGn7iVrga_d"
      },
      "outputs": [],
      "source": [
        "%%writefile prompts/create_dd.md\n",
        "\n",
        "<role>\n",
        "You are a world-class software architect specialized in Python.\n",
        "</role>\n",
        "\n",
        "<task>\n",
        "You will be given a Technical Specification located in the `technical_specification` tag below (\"the Spec\"). This Spec provides detailed requirements for designing one Python module.\n",
        "Your goal is to write a Technical Design Document (TDD) that suggests an optimal design for this Python module.\n",
        "</task>\n",
        "\n",
        "<instructions>\n",
        "- The design must strictly adhere to all requirements defined in the Spec, without adding any out-of-scope features.\n",
        "- The target runtime environment must be Python 3.11 on Linux.\n",
        "- The TDD must start with a brief overview of multiple viable design alternatives and compare their strengths and weaknesses.\n",
        "- The TDD must recommend a single design option and justify the choice with a clear trade-off analysis relative to the requirements.\n",
        "- The TDD must document the recommended design option in full details describing the system's structure and functionality using natural language, Mermaid diagrams, and short code snippets.\n",
        "- A Mermaid class diagram is recommended if the module design involves multiple related classes.\n",
        "- A Mermaid sequence diagram is recommended for non-trivial call chains or complex interactions.\n",
        "- The TDD must include a dedicated section outlining relevant security implications and proposed mitigations.\n",
        "</instructions>\n",
        "\n",
        "<response_format>\n",
        "- Write the TDD in Markdown format.\n",
        "</response_format>\n",
        "\n",
        "<technical_specification>\n",
        "{spec}\n",
        "</technical_specification>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZCoVBy1Vga_d"
      },
      "outputs": [],
      "source": [
        "%%writefile prompts/write_new_code.md\n",
        "\n",
        "<role>\n",
        "You are a world-class software engineer specialized in implementation of Python modules.\n",
        "</role>\n",
        "\n",
        "<task>\n",
        "You are given two documents:\n",
        "1. Software Requirements Specification located in the `software_requirements_specification` tag below.\n",
        "2. Technical Design Document located in the `technical_design_document` tag below.\n",
        "Your goal is to develop a new Python module based strictly on these two documents.\n",
        "</task>\n",
        "\n",
        "<instructions>\n",
        "- You can leverage all Standard Python Libraries available in Python 3.11, but you are not allowed to import any PyPi packages.\n",
        "- Before submitting your response perform a thorough code review and focus on the following 3 critical aspects:\n",
        "  1. Your code implementation must implement all requirements found in the specification and design doc, without any gaps.\n",
        "  2. Your code must be bug free.\n",
        "  3. Your code must not introduce any extra functionality that is not specifically requested in the design doc.\n",
        "</instructions>\n",
        "\n",
        "<response_format>\n",
        "Your response must be a single Markdown code block with the full code of the Python module that you implemented. Do not include any unit tests or code samples that show how to use this module.\n",
        "</response_format>\n",
        "\n",
        "<software_requirements_specification>\n",
        "{spec}\n",
        "</software_requirements_specification>\n",
        "\n",
        "<technical_design_document>\n",
        "{design_doc}\n",
        "</technical_design_document>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lBPtZNR0ga_d"
      },
      "source": [
        "### Step 1: Select and load Target Module\n",
        "\n",
        "We will use `twisted.python.zipstream` and download the raw file directly from GitHub.\n",
        "\n",
        "https://github.com/twisted/twisted/blob/master/src/twisted/python/zipstream.py\n",
        "\n",
        "*   **Source**: `original_code/zipstream.py`\n",
        "*   **Target**: `output/zipstream.py`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ODxstIc6ga_d"
      },
      "outputs": [],
      "source": [
        "MODULE_URL = \"https://raw.githubusercontent.com/twisted/twisted/master/src/twisted/python/zipstream.py\"\n",
        "TARGET_MODULE = \"zipstream\"\n",
        "CODE_DIR_PATH = \"original_code\"\n",
        "CODE_PATH = f\"{CODE_DIR_PATH}/{TARGET_MODULE}.py\"\n",
        "\n",
        "!wget -q -O {CODE_PATH} {MODULE_URL}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "59V3ymjRzSq-"
      },
      "source": [
        "### Step 2: Generate and Verify Unit Tests\n",
        "\n",
        "If the legacy code lacks tests, we must generate them to establish a baseline of correctness.\n",
        "\n",
        "#### 2.1 Generate Tests\n",
        "\n",
        "We read the legacy code and use a specific prompt to generate high-coverage `unittest` cases."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k5dYJPLRoMI3"
      },
      "outputs": [],
      "source": [
        "# Load code\n",
        "with open(CODE_PATH) as f:\n",
        "    source_code = f.read()\n",
        "print(f\"Loaded {TARGET_MODULE} ({len(source_code)} bytes)\")\n",
        "\n",
        "TEST_PATH = f\"output/tests/test_{TARGET_MODULE}.py\"\n",
        "\n",
        "generated_tests = generate_and_save(\n",
        "    template_path=\"prompts/create_tests.md\",\n",
        "    output_path=TEST_PATH,\n",
        "    replacements={\n",
        "        \"source_code\": source_code,\n",
        "        \"module_name\": TARGET_MODULE,\n",
        "        \"code_dir_path\": CODE_DIR_PATH,\n",
        "    },\n",
        ")\n",
        "print(f\"Tests saved to {TEST_PATH}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7U_0eSFUuMMz"
      },
      "source": [
        "#### 2.2 Verify Baseline (Run Tests)\n",
        "\n",
        "We must ensure these newly generated tests actually pass on the *legacy* code before moving on. We use `subprocess` to isolate the test execution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u_fabIYiuOhO"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import subprocess\n",
        "import sys\n",
        "\n",
        "from IPython.display import Markdown\n",
        "\n",
        "def run_tests(test_dir: str, code_dir: str, test_file: str) -> None:\n",
        "    \"\"\"Runs tests in a separate process and displays a rich Markdown table report.\"\"\"\n",
        "    env = os.environ.copy()\n",
        "    # Ensure the module under test can be imported\n",
        "    env[\"PYTHONPATH\"] = f\"{code_dir}\" + os.pathsep + env.get(\"PYTHONPATH\", \"\")\n",
        "\n",
        "    # Use -v to get detailed output for each test\n",
        "    cmd = [sys.executable, \"-m\", \"unittest\", \"-v\", os.path.join(test_dir, test_file)]\n",
        "\n",
        "    display(Markdown(f\"### Running `{test_file}` against `{code_dir}`...\"))\n",
        "    result = subprocess.run(cmd, check=False, env=env, capture_output=True, text=True)\n",
        "\n",
        "    table_lines = []\n",
        "    table_lines.append(\"| Test Name | Description | Status |\")\n",
        "    table_lines.append(\"| :--- | :--- | :---: |\")\n",
        "\n",
        "    lines = result.stderr.splitlines()\n",
        "    pending_test_name = None\n",
        "    passed = 0\n",
        "    failed = 0\n",
        "    skipped = 0\n",
        "\n",
        "    for line in lines:\n",
        "        line = line.strip()\n",
        "        if not line:\n",
        "            continue\n",
        "\n",
        "        if \" ... \" in line:\n",
        "            parts = line.split(\" ... \")\n",
        "            if len(parts) >= 2:\n",
        "                desc_part = parts[0].strip()\n",
        "                status_part = parts[1].strip()\n",
        "\n",
        "                if pending_test_name:\n",
        "                    test_name = pending_test_name\n",
        "                    description = desc_part\n",
        "                    pending_test_name = None\n",
        "                else:\n",
        "                    # Fallback parse if no header line preceded\n",
        "                    test_parts = desc_part.split(\" \")\n",
        "                    test_name = test_parts[0]\n",
        "                    description = (\n",
        "                        \" \".join(test_parts[1:]).strip() if len(test_parts) > 1 else \"\"\n",
        "                    )\n",
        "                    # Handle trailing parens if present\n",
        "                    if \"(\" in description:\n",
        "                        description = description.split(\"(\")[0].strip()\n",
        "\n",
        "                if status_part == \"ok\":\n",
        "                    status_emoji = \"✅ PASS\"\n",
        "                    passed += 1\n",
        "                elif status_part.startswith(\"FAIL\"):\n",
        "                    status_emoji = \"❌ FAIL\"\n",
        "                    failed += 1\n",
        "                elif status_part.startswith(\"ERROR\"):\n",
        "                    status_emoji = \"🛑 ERROR\"\n",
        "                    failed += 1\n",
        "                elif status_part.startswith(\"skipped\"):\n",
        "                    status_emoji = \"⚠️ SKIP\"\n",
        "                    skipped += 1\n",
        "                else:\n",
        "                    status_emoji = f\"❓ {status_part.split()[0].upper()}\"\n",
        "\n",
        "                table_lines.append(\n",
        "                    f\"| `{test_name}` | {description} | {status_emoji} |\"\n",
        "                )\n",
        "\n",
        "        elif line.startswith(\"test_\") and \"(\" in line and \")\" in line:\n",
        "            pending_test_name = line.split(\" \")[0]\n",
        "        elif not line.startswith(\"test_\") and line.strip() != \"\":\n",
        "            pending_test_name = None\n",
        "\n",
        "    display(Markdown(\"\\n\".join(table_lines)))\n",
        "    display(\n",
        "        Markdown(f\"\\n**Summary:** {passed} passed, {failed} failed, {skipped} skipped.\")\n",
        "    )\n",
        "\n",
        "    if result.returncode == 0 and failed == 0:\n",
        "        display(Markdown(\"\\n### ✅ **OVERALL SUCCESS: All tests passed.**\"))\n",
        "        return True\n",
        "    display(Markdown(\"\\n### ❌ **OVERALL FAILURE: Some tests failed.**\"))\n",
        "    if failed > 0 or result.returncode != 0:\n",
        "        print(\"\\nDetailed error output:\")\n",
        "        print(result.stderr)\n",
        "    return False\n",
        "\n",
        "\n",
        "def verify_module(test_dir: str, code_dir: str, module_name: str) -> bool:\n",
        "    \"\"\"Wrapper to check file existence and run tests.\"\"\"\n",
        "    code_path = os.path.join(code_dir, f\"{module_name}.py\")\n",
        "    test_file = f\"test_{module_name}.py\"\n",
        "    if os.path.exists(code_path) and os.path.exists(os.path.join(test_dir, test_file)):\n",
        "        success = run_tests(test_dir, code_dir, test_file)\n",
        "        if success:\n",
        "            print(f\"\\nVerified! Valid module found at: {code_path}\")\n",
        "        return success\n",
        "    print(\"Cannot run verification due to missing files.\")\n",
        "    return False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8JWqlqSGg8qK"
      },
      "outputs": [],
      "source": [
        "# Run tests on original code\n",
        "verify_module(\"output/tests\", CODE_DIR_PATH, TARGET_MODULE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eJPdp09luTJ9"
      },
      "source": [
        "### Step 3: Generate Specification\n",
        "\n",
        "Create a detailed Markdown specification derived strictly from the legacy code's behavior."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "suCuBQy0uXRO"
      },
      "outputs": [],
      "source": [
        "SPEC_PATH = f\"output/specs/{TARGET_MODULE}.md\"\n",
        "spec_content = generate_and_save(\n",
        "    template_path=\"prompts/create_spec.md\",\n",
        "    output_path=SPEC_PATH,\n",
        "    replacements={\"source_code\": source_code},\n",
        ")\n",
        "display(Markdown(spec_content))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s4h4KxrYzVY_"
      },
      "source": [
        "### Step 4: Generate Design Document\n",
        "\n",
        "Using *only* the specification from Step 3, generate a design document. This ensures the new design is based on desired behavior, not just copying the old implementation details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5cm70qIBga_e"
      },
      "outputs": [],
      "source": [
        "DD_PATH = f\"output/design_docs/{TARGET_MODULE}.md\"\n",
        "dd_content = generate_and_save(\n",
        "    template_path=\"prompts/create_dd.md\",\n",
        "    output_path=DD_PATH,\n",
        "    replacements={\"spec\": spec_content},\n",
        ")\n",
        "render_markdown_with_mermaid(dd_content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QJABk09Yz1s5"
      },
      "source": [
        "### Step 5: Generate New Code\n",
        "\n",
        "Generate the refactored code using the Specification and Design Doc. The legacy code is *not* used in this input, preventing contamination from old anti-patterns."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-RnB1sdnga_e"
      },
      "outputs": [],
      "source": [
        "NEW_CODE_PATH = f\"output/code/{TARGET_MODULE}.py\"\n",
        "new_code = generate_and_save(\n",
        "    template_path=\"prompts/write_new_code.md\",\n",
        "    output_path=NEW_CODE_PATH,\n",
        "    replacements={\"spec\": spec_content, \"design_doc\": dd_content},\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "50PZ7q9v0J7J"
      },
      "source": [
        "### Step 6: Final Verification\n",
        "\n",
        "Run the tests generated in Step 2 against the new code generated in Step 5."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1W5MjFSz0Ljj"
      },
      "outputs": [],
      "source": [
        "verify_module(\"output/tests\", \"output/code\", TARGET_MODULE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yf-gr2S40M3o"
      },
      "source": [
        "## 7. Introduce new Requirements for the code\n",
        "\n",
        "Delete and rewrite the code, show that the added requirements are successfully incorporated.\n",
        "\n",
        "We will update the specification with new requirements, then re-generate the Design Doc and code based upon that specification."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hwxPIKBFga_f"
      },
      "outputs": [],
      "source": [
        "# Add new requirements to specification\n",
        "new_requirements = \"\"\"## Coding style requirements\n",
        "1. All input parameters and return values must have type hints.\n",
        "2. All methods and classes must have Google-style Docstrings.\n",
        "\"\"\"\n",
        "\n",
        "with open(SPEC_PATH, \"a+\") as f:\n",
        "    f.write(new_requirements)\n",
        "    f.seek(0)\n",
        "    spec_content = f.read()\n",
        "\n",
        "# Regenerate Design Doc & Code\n",
        "dd_content = generate_and_save(\"prompts/create_dd.md\", DD_PATH, {\"spec\": spec_content})\n",
        "new_code = generate_and_save(\n",
        "    \"prompts/write_new_code.md\",\n",
        "    NEW_CODE_PATH,\n",
        "    {\"spec\": spec_content, \"design_doc\": dd_content},\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HvnpOMFHa6or"
      },
      "outputs": [],
      "source": [
        "verify_module(\"output/tests\", \"output/code\", TARGET_MODULE)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "refactor_legacy_code.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
