Added optimizer

src/codergpt/__init__.py

@@ -4,6 +4,7 @@
 
 from codergpt.commenter import CodeCommenter
 from codergpt.explainer import CodeExplainer
+from codergpt.optimizer import CodeOptimizer
 
 from .main import CoderGPT
 
@@ -14,4 +15,4 @@
     __version__ = "0.0.0"  # pragma: no cover
 
 
-__all__ = ["CoderGPT", "CodeExplainer", "CodeCommenter"]
+__all__ = ["CoderGPT", "CodeExplainer", "CodeCommenter", "CodeOptimizer"]

src/codergpt/cli.py

@@ -36,6 +36,8 @@
 logger = logging.getLogger(__name__)
 
 path_argument = click.argument("path", type=click.Path(exists=True))
+function_option = click.option("-f", "--function", help="Function name to explain or optimize.")
+class_option = click.option("-c", "--classname", help="Class name to explain or optimize.")
 overwrite_option = click.option(
     "--overwrite/--no-overwrite", is_flag=True, default=False, help="Overwrite the existing file."
 )
@@ -77,8 +79,8 @@ def inspect(path: Union[str, Path, TextIO]):
 
 @main.command()
 @path_argument
-@click.option("-f", "--function", help="Function name to explain.")
-@click.option("-c", "--classname", help="Class name to explain.")
+@function_option
+@class_option
 def explain(path: Union[str, Path], function: str, classname: str):
     """
     Inspect package to show file-language-map.
@@ -119,5 +121,30 @@ def add_comments(path: Union[str, Path], overwrite: bool = False):
         raise ValueError("The path provided is not a file.")
 
 
+@main.command("optimize")
+@path_argument
+@function_option
+@class_option
+@overwrite_option
+def optimize_code(path: Union[str, Path], function: str, classname: str, overwrite: bool = False):
+    """
+    Optimize the code file.
+
+    :param path: The path to the code file.
+    :param function: The name of the function to optimize. Default is None.
+    :param classname: The name of the class to optimize. Default is None.
+    :param overwrite: Whether to overwrite the existing file. Default is False.
+    """
+    # Ensure path is a string or Path object for consistency
+    if isinstance(path, str):
+        path = Path(path)
+
+    # Check if path is a file
+    if path.is_file():
+        coder.optimizer(path=path, function=function, classname=classname, overwrite=overwrite)
+    else:
+        raise ValueError("The path provided is not a file.")
+
+
 if __name__ == "__main__":
     main()

src/codergpt/commenter/commenter.py

@@ -28,7 +28,7 @@ def comment(self, code: str, filename: str, overwrite: bool = False, language: O
         response = self.chain.invoke(
             {
                 "input": f"Rewrite and return this {language} code with\
-                comments and docstrings in :param: format: \n{code}\n"
+                comments and sphinx docstrings in :param: format: \n{code}\n"
             }
         )
 

src/codergpt/main.py

@@ -12,6 +12,7 @@
 from codergpt.commenter.commenter import CodeCommenter
 from codergpt.constants import EXTENSION_MAP_FILE, INSPECTION_HEADERS
 from codergpt.explainer.explainer import CodeExplainer
+from codergpt.optimizer.optimizer import CodeOptimizer
 
 
 class CoderGPT:
@@ -115,6 +116,24 @@ def commenter(self, path: Union[str, Path], overwrite: bool = False):
         code, language = self.get_code(filename=path)
         code_commenter.comment(code=code, filename=path, overwrite=overwrite, language=language)
 
+    def optimizer(self, path: Union[str, Path], function: str = None, classname=None, overwrite: bool = False):
+        """
+        Optimize the code file.
+
+        :param path: The path to the code file.
+        """
+        code_optimizer = CodeOptimizer(self.chain)
+        # code, language = self.get_code(filename=path, function_name=function, class_name=classname)
+        code_optimizer.optimize(filename=path, function=function, classname=classname, overwrite=overwrite)
+
+    def tester(self, path: Union[str, Path]):
+        """
+        Test the code file.
+
+        :param path: The path to the code file.
+        """
+        pass
+
 
 if __name__ == "__main__":
     coder = CoderGPT()

src/codergpt/optimizer/__init__.py

@@ -1 +1,5 @@
 """Code optimization module."""
+
+from .optimizer import CodeOptimizer
+
+__all__ = ["CodeOptimizer"]

src/codergpt/optimizer/optimizer.py

@@ -0,0 +1,70 @@
+"""Optimizer Module."""
+
+import os
+from typing import Any, Dict, Optional
+
+from langchain_core.runnables.base import RunnableSerializable
+
+
+class CodeOptimizer:
+    """Code Explainer class that extracts and explains code from a given file."""
+
+    def __init__(self, chain: RunnableSerializable[Dict, Any]):
+        """
+        Initialize the CodeExplainer class with a runnable chain.
+
+        :param chain: A RunnableSerializable object capable of executing tasks.
+        """
+        self.chain = chain
+
+    def optimize(
+        self,
+        filename: str,
+        function: Optional[str] = None,
+        classname: Optional[str] = None,
+        overwrite: bool = False,
+    ):
+        """
+        Optimize the the code by invoking the runnable chain.
+
+        :param path: The path to the code file to be explained.
+        :param function: The name of the function to explain. Default is None.
+        :param classname: The name of the class to explain. Default is None.
+        """
+        with open(filename, "r") as source_file:
+            source_code = source_file.read()
+        if function:
+            response = self.chain.invoke(
+                {
+                    "input": f"Optimize, comment and add sphinx docstrings"
+                    f" to the function '{function}' in \n\n```\n{source_code}\n```"
+                    "Also explain the optimization ina systematic way."
+                }
+            )
+        elif classname:
+            response = self.chain.invoke(
+                {
+                    "input": f"Optimize, comment and add sphinx docstrings"
+                    f" to the class '{classname}' in \n\n```\n{source_code}\n```"
+                    "Also explain the optimization ina systematic way."
+                }
+            )
+        else:
+            # Explain full code
+            response = self.chain.invoke(
+                {
+                    "input": f"Optimize, comment and add sphinx style docstrings"
+                    f" to the following code: \n\n```\n{source_code}\n```"
+                    "Also explain the optimization ina systematic way."
+                }
+            )
+        optimized_code = response.content
+        new_filename = filename
+        if not overwrite:
+            # Create a new filename with the _updated suffix
+            base, ext = os.path.splitext(filename)
+            new_filename = f"{base}_updated{ext}"
+
+        # Write the commented code to the new file
+        with open(new_filename, "w") as updated_file:
+            updated_file.write(optimized_code)

tests/test_optimizer.py

@@ -0,0 +1,148 @@
+"""Optimizer tests."""
+
+import unittest
+from unittest.mock import MagicMock, mock_open, patch
+
+from codergpt import CodeOptimizer
+
+from tests.test_constants import TEST_INPUT_DIR, TEST_OUTPUT_DIR
+
+
+class TestCodeOptimizer(unittest.TestCase):
+    """Test the CodeOptimizer class."""
+
+    def setUp(self):
+        """
+        Set up the test case.
+
+        Done by creating a mock chain that can be used to simulate the behavior of the actual chain.
+        """
+        self.mock_chain = MagicMock()
+        self.code_optimizer = CodeOptimizer(chain=self.mock_chain)
+        self.bad_python = TEST_INPUT_DIR / "bad_python.py"
+        self.expected_bad_python = TEST_OUTPUT_DIR / "expected_bad_python.py"
+
+    @patch("builtins.open", new_callable=mock_open, read_data="def foo(): pass")
+    def test_optimize_with_function(self, mock_file):
+        """Test the optimize method with a function."""
+        # Setup
+        filename = "test.py"
+        function = "foo"
+        expected_response = MagicMock(content="def foo():\n    # Optimized code\n    pass")
+        self.mock_chain.invoke.return_value = expected_response
+
+        # Exercise
+        self.code_optimizer.optimize(filename, function=function)
+
+        # Verify
+        self.mock_chain.invoke.assert_called_once()
+        mock_file.assert_called_with(filename.replace(".py", "_updated.py"), "w")
+        mock_file().write.assert_called_once_with(expected_response.content)
+
+    @patch("builtins.open", new_callable=mock_open, read_data="class Bar: pass")
+    def test_optimize_with_classname(self, mock_file):
+        """
+        Test the optimize method with a classname.
+
+        :param self: The current instance of the class
+        """
+        # Setup
+        filename = "test.py"
+        classname = "Bar"
+        expected_response = MagicMock(content="class Bar:\n    # Optimized code\n    pass")
+        self.mock_chain.invoke.return_value = expected_response
+
+        # Exercise
+        self.code_optimizer.optimize(filename, classname=classname)
+
+        # Verify
+        self.mock_chain.invoke.assert_called_once()
+        mock_file.assert_called_with(filename.replace(".py", "_updated.py"), "w")
+        mock_file().write.assert_called_once_with(expected_response.content)
+
+    @patch("builtins.open", new_callable=mock_open, read_data="def foo(): pass\nclass Bar: pass")
+    def test_optimize_full_code(self, mock_file):
+        """
+        Test the optimize method with full code.
+
+        :param self: The current instance of the class
+        """
+        # Setup
+        filename = "test.py"
+        expected_response = MagicMock(
+            content="def foo():\n    # Optimized function\n    pass\nclass Bar:\n    # Optimized class\n    pass"
+        )
+        self.mock_chain.invoke.return_value = expected_response
+
+        # Exercise
+        self.code_optimizer.optimize(filename)
+
+        # Verify
+        self.mock_chain.invoke.assert_called_once()
+        mock_file.assert_called_with(filename.replace(".py", "_updated.py"), "w")
+        mock_file().write.assert_called_once_with(expected_response.content)
+
+    @patch("builtins.open", new_callable=mock_open, read_data="def foo(): pass")
+    def test_optimize_overwrite_false(self, mock_file):
+        """
+        Test that when overwrite is True, the original file is overwritten.
+
+        :param self: The current instance of the class
+        """
+        """
+        Test that when overwrite is True, the original file is overwritten.
+        """
+        # Setup
+        filename = "test.py"
+        expected_response = MagicMock(content="def foo():\n    # Optimized function\n    pass")
+        self.mock_chain.invoke.return_value = expected_response
+
+        # Exercise
+        self.code_optimizer.optimize(filename)
+
+        # Verify
+        self.mock_chain.invoke.assert_called_once()
+        mock_file.assert_called_with(filename.replace(".py", "_updated.py"), "w")
+        mock_file().write.assert_called_once_with(expected_response.content)
+
+    @patch("builtins.open", new_callable=mock_open, read_data="def foo(): pass")
+    def test_optimize_overwrite_true(self, mock_file):
+        """
+        Test that when an invalid language is provided, the appropriate error is raised or handled.
+
+        :param self: The current instance of the class
+        """
+        # Setup
+        filename = "test.py"
+        expected_response = MagicMock(content="def foo():\n    # Optimized function\n    pass")
+        self.mock_chain.invoke.return_value = expected_response
+
+        # Exercise
+        self.code_optimizer.optimize(filename, overwrite=True)
+
+        # Verify
+        self.mock_chain.invoke.assert_called_once()
+        mock_file.assert_called_with(filename, "w")
+        mock_file().write.assert_called_once_with(expected_response.content)
+
+    #  TODO: Fix this test
+    # @patch("builtins.open", new_callable=mock_open, read_data="def foo(): pass")
+    # def test_optimize_invalid_function(self, mock_file):
+    #     """
+    #     Test that when an invalid function is provided, the appropriate error is raised or handled.
+
+    #     :param self: The current instance of the class
+    #     """
+    #     # Setup
+    #     filename = "test.py"
+    #     function = "invalid_function"
+    #     expected_error_message = "Function not found in the source code."
+
+    #     # Exercise & Verify
+    #     with self.assertRaises(ValueError) as context:
+    #         self.code_optimizer.optimize(filename, function=function)
+    #     self.assertEqual(str(context.exception), expected_error_message)
+
+
+if __name__ == "__main__":
+    unittest.main()

tox.ini

@@ -82,5 +82,5 @@ skip_install = true
 deps =
     docstr-coverage
 commands =
-    docstr-coverage src/ tests/ --skip-private --skip-magic --exclude=tests/input/* --exclude=tests/output/*
+    docstr-coverage src/ tests/ --skip-private --skip-magic --exclude .*tests/(input|output)/*
 description = Run the docstr-coverage tool to check documentation coverage.