Added documentation

README.md

@@ -192,6 +192,50 @@ code [OPTIONS] COMMAND [ARGS]...
     """
     ```
 
+5. `write-tests`: Writes tests for the specified code file. The user can specify a function and/or a class within the file to target with the tests.
+
+```shell
+code write-tests <path> [--function <function_name>] [--class <classname>]
+```
+
+#### Example
+- Let's consider a python file `example.py`:
+```python
+# example.py
+
+def add(a, b):
+    return a + b
+
+class Calculator:
+    def subtract(self, a, b):
+        return a - b
+```
+```shell
+$ code write-tests example.py --function add --class Calculator
+```
+results in test files being generated that contain test cases for the `add` function and the `Calculator` class. The actual content of the test files will depend on the implementation of the `coder.test_writer` method but would typically look something like this:
+
+```python
+import unittest
+from example import add, Calculator
+
+class TestAddFunction(unittest.TestCase):
+
+    def test_addition(self):
+        self.assertEqual(add(3, 4), 7)
+
+class TestCalculator(unittest.TestCase):
+
+    def setUp(self):
+        self.calc = Calculator()
+
+    def test_subtract(self):
+        self.assertEqual(self.calc.subtract(10, 5), 5)
+```
+
+In this example, running the command generates unit tests for both the `add` function and the `Calculator` class in the `example.py` file. The tests check if the `add` function correctly adds two numbers and if the `Calculator`'s `subtract` method correctly subtracts one number from another.
+
+
 ## Development
 
 The CLI is built using Python and the `click` library. Below is an example of how to define a new command:

docs/description.rst

@@ -202,6 +202,52 @@ Commands
         By using these optimizations, we improve the efficiency and readability of the code.
         """
 
+5s. **write-tests**: Generates test cases for specified functions and/or classes within a Python code file.
+
+   .. code-block:: shell
+
+       code write-tests <path> [--function <function_name>] [--class <classname>]
+
+   **Example**
+
+   - Let's consider a Python file `example.py`:
+
+     .. code-block:: python
+
+        # example.py
+
+        def add(a, b):
+            return a + b
+
+        class Calculator:
+            def subtract(self, a, b):
+                return a - b
+
+   .. code-block:: shell
+
+       $ code write-tests example.py --function add --class Calculator
+
+   results in the creation of test files that contain test cases for both the `add` function and the `Calculator` class. The content of the generated test files might look like this:
+
+   .. code-block:: python
+
+        import unittest
+        from example import add, Calculator
+
+        class TestAddFunction(unittest.TestCase):
+
+            def test_addition(self):
+                self.assertEqual(add(3, 4), 7)
+
+        class TestCalculator(unittest.TestCase):
+
+            def setUp(self):
+                self.calc = Calculator()
+
+            def test_subtract(self):
+                self.assertEqual(self.calc.subtract(10, 5), 5)
+
+   In this example, executing the command generates unit tests for the `add` function and the `Calculator` class defined in `example.py`. The tests verify whether the `add` function correctly computes the sum of two numbers and if the `Calculator`'s `subtract` method accurately performs subtraction.
 
 Development
 -----------

docs/test.rst

@@ -0,0 +1,30 @@
+.. py:module:: codergpt
+
+Test writing module
+===================
+
+.. py:class:: CodeTester(chain)
+
+   The CodeTester class is responsible for generating testing code from a given source file. It utilizes a llm chain to produce tests for specific functions or classes within the source file.
+
+   .. py:method:: __init__(chain)
+
+      Initializes the CodeTester instance with a provided llm chain.
+
+      :param chain: A RunnableSerializable object capable of executing tasks.
+      :type chain: RunnableSerializable[Dict, Any]
+
+   .. py:method:: write_tests(filename, function=None, classname=None, outfile=None)
+
+      Generates test cases for the specified code by invoking the llm chain. If a function or class name is provided, it will generate tests specifically for that function or class. Otherwise, it will attempt to create tests for the entire code.
+
+      :param filename: The path to the code file for which tests are to be written.
+      :type filename: Union[str, Path]
+      :param function: The name of the function for which tests should be generated. Defaults to None, indicating that no specific function is targeted.
+      :type function: Optional[str]
+      :param classname: The name of the class for which tests should be generated. Defaults to None, indicating that no specific class is targeted.
+      :type classname: Optional[str]
+      :param outfile: The path where the generated test file should be saved. If not provided, a default path within the TEST_DIR will be used.
+      :type outfile: Optional[str]
+
+      The method reads the source code from the provided filename and uses the llm chain to generate appropriate test cases. The resulting test code is then written to either the specified outfile or a new file within the TEST_DIR directory.

src/codergpt/optimizer/optimizer.py

@@ -33,22 +33,23 @@ def optimize(
         """
         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 in a systematic way as a comment."
-                }
-            )
-        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 in a systematic way."
-                }
-            )
+        if function or classname:
+            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 in a systematic way as a comment."
+                    }
+                )
+            if 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 in a systematic way."
+                    }
+                )
         else:
             # Optimize full code
             response = self.chain.invoke(

src/codergpt/test_writer/test_writer.py

@@ -36,20 +36,21 @@ def write_tests(
         """
         with open(filename, "r") as source_file:
             source_code = source_file.read()
-        if function:
-            response = self.chain.invoke(
-                {
-                    "input": f"Write tests for the function '{function}' in \n\n```\n{source_code}\n```"
-                    "Return just the code block. Also explain the tests in a systematic way as a comment."
-                }
-            )
-        elif classname:
-            response = self.chain.invoke(
-                {
-                    "input": f"Write tests for the class '{classname}' in \n\n```\n{source_code}\n```"
-                    "Also explain the tests in a systematic way."
-                }
-            )
+        if function or classname:
+            if function:
+                response = self.chain.invoke(
+                    {
+                        "input": f"Write tests for the function '{function}' in \n\n```\n{source_code}\n```"
+                        "Return just the code block. Also explain the tests in a systematic way as a comment."
+                    }
+                )
+            if classname:
+                response = self.chain.invoke(
+                    {
+                        "input": f"Write tests for the class '{classname}' in \n\n```\n{source_code}\n```"
+                        "Also explain the tests in a systematic way."
+                    }
+                )
         else:
             # Write tests for full code
             response = self.chain.invoke(