From d6e5ad3da2df133393d1598cdc5c65bb55efcb61 Mon Sep 17 00:00:00 2001
From: Harshad Hegde <hegdehb@gmail.com>
Date: Mon, 12 Feb 2024 16:38:03 -0600
Subject: [PATCH] added outfile

---
 src/codergpt/cli.py                 | 11 +++++++----
 src/codergpt/commenter/commenter.py |  1 +
 src/codergpt/main.py                |  8 ++++----
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/codergpt/cli.py b/src/codergpt/cli.py
index e6af281..8760569 100644
--- a/src/codergpt/cli.py
+++ b/src/codergpt/cli.py
@@ -41,6 +41,7 @@
 overwrite_option = click.option(
     "--overwrite/--no-overwrite", is_flag=True, default=False, help="Overwrite the existing file."
 )
+output_option = click.option("-o", "--outfile", help="Output file to write to.")
 
 coder = CoderGPT()
 
@@ -150,7 +151,8 @@ def optimize_code(path: Union[str, Path], function: str, classname: str, overwri
 @path_argument
 @function_option
 @class_option
-def write_test_code(path: Union[str, Path], function: str, classname: str):
+@output_option
+def write_test_code(path: Union[str, Path], function: str, classname: str, outfile: Union[str, Path] = None):
     """
     Write tests for the code file.
 
@@ -164,14 +166,15 @@ def write_test_code(path: Union[str, Path], function: str, classname: str):
 
     # Check if path is a file
     if path.is_file():
-        coder.test_writer(path=path, function=function, classname=classname)
+        coder.test_writer(path=path, function=function, classname=classname, outfile=outfile)
     else:
         raise ValueError("The path provided is not a file.")
 
 
 @main.command("document")
 @path_argument
-def write_documentation(path: Union[str, Path]):
+@output_option
+def write_documentation(path: Union[str, Path], outfile: Union[str, Path] = None):
     """
     Write documentation files for the code file.
 
@@ -183,7 +186,7 @@ def write_documentation(path: Union[str, Path]):
 
     # Check if path is a file
     if path.is_file():
-        coder.documenter(path=path)
+        coder.documenter(path=path, outfile=outfile)
     else:
         raise ValueError("The path provided is not a file.")
 
diff --git a/src/codergpt/commenter/commenter.py b/src/codergpt/commenter/commenter.py
index 65304b5..87ebac9 100644
--- a/src/codergpt/commenter/commenter.py
+++ b/src/codergpt/commenter/commenter.py
@@ -29,6 +29,7 @@ def comment(self, code: str, filename: str, overwrite: bool = False, language: O
             {
                 "input": f"Rewrite and return this {language} code with\
                 comments and sphinx docstrings in :param: format: \n{code}\n"
+                "Return just the code block since all this will be a file."
             }
         )
 
diff --git a/src/codergpt/main.py b/src/codergpt/main.py
index f8e577b..eeb0017 100644
--- a/src/codergpt/main.py
+++ b/src/codergpt/main.py
@@ -128,23 +128,23 @@ def optimizer(self, path: Union[str, Path], function: str = None, classname=None
         # 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 test_writer(self, path: Union[str, Path], function: str = None, classname: str = None):
+    def test_writer(self, path: Union[str, Path], function: str = None, classname: str = None, outfile: str = None):
         """
         Test the code file.
 
         :param path: The path to the code file.
         """
         code_tester = CodeTester(self.chain)
-        code_tester.write_tests(filename=path, function=function, classname=classname)
+        code_tester.write_tests(filename=path, function=function, classname=classname, outfile=outfile)
 
-    def documenter(self, path: Union[str, Path]):
+    def documenter(self, path: Union[str, Path], outfile: str = None):
         """
         Document the code file.
 
         :param path: The path to the code file.
         """
         code_documenter = CodeDocumenter(self.chain)
-        code_documenter.document(filename=path)
+        code_documenter.document(filename=path, outfile=outfile)
 
 
 if __name__ == "__main__":