Merge pull request #40 from MichaelisTrofficus/add-exception-catching

Add exception catching

Author: MichaelisTrofficus

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "gpt4docstrings"
-version = "0.1.0"
+version = "0.1.1"
 description = "gpt4docstrings"
 authors = ["Miguel Otero Pedrido <miguel.otero.pedrido.1993@gmail.com>"]
 license = "MIT"

src/gpt4docstrings/docstrings_generators/chatgpt_generator.py

@@ -11,6 +11,7 @@
 from gpt4docstrings.docstrings_generators.utils.parsers import DocstringParser
 from gpt4docstrings.docstrings_generators.utils.prompts import CLASS_PROMPTS
 from gpt4docstrings.docstrings_generators.utils.prompts import FUNCTION_PROMPTS
+from gpt4docstrings.exceptions import ASTError
 
 
 class ChatGPTDocstringGenerator:
@@ -57,6 +58,9 @@ def generate_function_docstring(self, source: str) -> dict:
 
         Returns:
             dict: A dictionary containing the generated docstring.
+
+        Raises:
+            ASTError: Raises an ASTError when there are errors interacting with an AST node
         """
         source = source.strip()
         stripped_source = textwrap.dedent(source)
@@ -66,14 +70,19 @@ def generate_function_docstring(self, source: str) -> dict:
         )
         _input = prompt.format_prompt(code=stripped_source)
         fn_src = DocstringParser().parse(self._get_completion(_input.to_string()))
-        fn_node = RedBaron(fn_src)[0]
 
-        return {
-            "docstring": utils.add_indentation_to_docstring(
-                '"""' + textwrap.dedent(fn_node[0].to_python()) + '"""',
-                fn_node[0].indentation,
-            )
-        }
+        try:
+            fn_node = RedBaron(fn_src)[0]
+            return {
+                "docstring": utils.add_indentation_to_docstring(
+                    '"""' + textwrap.dedent(fn_node[0].to_python()) + '"""',
+                    fn_node[0].indentation,
+                )
+            }
+        except ValueError as e:
+            raise ASTError(
+                f"Some error has occurred when trying to parse the current AST node: {e}"
+            ) from e
 
     def generate_class_docstring(self, source: str) -> dict:
         """
@@ -84,6 +93,9 @@ def generate_class_docstring(self, source: str) -> dict:
 
         Returns:
             dict: A dictionary containing the generated docstrings.
+
+        Raises:
+            ASTError: Raises an ASTError when there are errors interacting with an AST node
         """
         source = source.strip()
         stripped_source = textwrap.dedent(source)
@@ -93,20 +105,27 @@ def generate_class_docstring(self, source: str) -> dict:
         )
         _input = prompt.format_prompt(code=stripped_source)
         class_src = DocstringParser().parse(self._get_completion(_input.to_string()))
-        class_node = RedBaron(class_src)[0]
-        method_nodes = [f for f in class_node.find_all("def")]
-
-        docstrings = {}
-        for method_node in method_nodes:
-            docstrings[method_node.name] = utils.add_indentation_to_docstring(
-                '"""' + textwrap.dedent(method_node[0].to_python()) + '"""',
-                method_node[0].indentation,
+
+        try:
+            class_node = RedBaron(class_src)[0]
+            method_nodes = [f for f in class_node.find_all("def")]
+
+            docstrings = {}
+            for method_node in method_nodes:
+                docstrings[method_node.name] = utils.add_indentation_to_docstring(
+                    '"""' + textwrap.dedent(method_node[0].to_python()) + '"""',
+                    method_node[0].indentation,
+                )
+
+            docstrings["docstring"] = class_node.value[0]
+            docstrings["docstring"] = utils.add_indentation_to_docstring(
+                '"""' + textwrap.dedent(class_node[0].to_python()) + '"""',
+                class_node[0].indentation,
             )
 
-        docstrings["docstring"] = class_node.value[0]
-        docstrings["docstring"] = utils.add_indentation_to_docstring(
-            '"""' + textwrap.dedent(class_node[0].to_python()) + '"""',
-            class_node[0].indentation,
-        )
+            return docstrings
 
-        return docstrings
+        except ValueError as e:
+            raise ASTError(
+                f"Some error has occurred when trying to parse the current AST node: {e}"
+            ) from e

src/gpt4docstrings/docstrings_generators/utils/parsers.py

@@ -2,11 +2,7 @@
 
 from langchain.schema import BaseOutputParser
 
-
-class DocstringParsingError(Exception):
-    """Custom exception for docstring parsing errors."""
-
-    pass
+from gpt4docstrings.exceptions import DocstringParsingError
 
 
 class DocstringParser(BaseOutputParser):

src/gpt4docstrings/exceptions.py

@@ -0,0 +1,10 @@
+class DocstringParsingError(Exception):
+    """Custom exception for docstring parsing errors."""
+
+    pass
+
+
+class ASTError(Exception):
+    """Custom exception for AST related error"""
+
+    pass

src/gpt4docstrings/generate_docstrings.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import pathlib
 import sys
@@ -13,6 +14,8 @@
 from gpt4docstrings import utils
 from gpt4docstrings.ascii_title import title
 from gpt4docstrings.docstrings_generators import ChatGPTDocstringGenerator
+from gpt4docstrings.exceptions import ASTError
+from gpt4docstrings.exceptions import DocstringParsingError
 
 
 class GPT4Docstrings:
@@ -126,6 +129,7 @@ def get_filenames_from_paths(self) -> List[str]:
         self.common_base = utils.get_common_base(filenames)
         return filenames
 
+    # flake8: noqa: C901
     def generate_file_docstrings(self, filename: str):
         """Generates docstrings for a single file.
 
@@ -136,33 +140,58 @@ def generate_file_docstrings(self, filename: str):
         click.echo(click.style(f"Documenting file {filename} ... ", fg="green"))
 
         for node in source.find_all("def"):
-            if not utils.check_def_node_is_class_method(node):
+            if not utils.check_def_node_is_class_method(
+                node
+            ) and not utils.check_def_node_is_nested(node):
                 if not node.value[0].type == "string":
-                    fn_docstring = self.docstring_generator.generate_function_docstring(
-                        node.dumps()
-                    )
-                    node.value.insert(0, fn_docstring["docstring"])
-                    self.documented_nodes.append([filename, node.name])
+                    try:
+                        fn_docstring = (
+                            self.docstring_generator.generate_function_docstring(
+                                node.dumps()
+                            )
+                        )
+                        node.value.insert(0, fn_docstring["docstring"])
+                        self.documented_nodes.append([filename, node.name])
+                    except DocstringParsingError:
+                        logging.warning(
+                            f"Skipping {node.name} from {filename} due to errors when parsing"
+                        )
+                    except ASTError:
+                        logging.warning(
+                            f"Skipping {node.name} from {filename} due to errors when accessing AST node"
+                        )
 
         for node in source.find_all("class"):
             if not node.value[0].type == "string":
-                class_docstring = self.docstring_generator.generate_class_docstring(
-                    node.dumps()
-                )
-                node.value.insert(
-                    0,
-                    class_docstring["docstring"],
-                )
-
-                for method_node in node.value:
-                    if (
-                        method_node.type == "def"
-                        and not utils.check_is_private_method(method_node)
-                        and not method_node.value[0].type == "string"
-                    ):
-                        method_node.value.insert(0, class_docstring[method_node.name])
-
-                self.documented_nodes.append([filename, node.name])
+                try:
+                    class_docstring = self.docstring_generator.generate_class_docstring(
+                        node.dumps()
+                    )
+                    node.value.insert(
+                        0,
+                        class_docstring["docstring"],
+                    )
+
+                    for method_node in node.value:
+                        if (
+                            method_node.type == "def"
+                            and not utils.check_is_private_method(method_node)
+                            and not method_node.value[0].type == "string"
+                            and not utils.check_def_node_is_nested(method_node)
+                        ):
+                            method_node.value.insert(
+                                0, class_docstring[method_node.name]
+                            )
+
+                    self.documented_nodes.append([filename, node.name])
+                except DocstringParsingError:
+                    logging.warning(
+                        f"Skipping {node.name} from {filename} due to errors when parsing"
+                    )
+                except ASTError:
+                    logging.warning(
+                        f"Skipping {node.name} from {filename} due to errors when accessing AST node"
+                    )
 
         utils.write_updated_source_to_file(source, filename)
 

src/gpt4docstrings/utils.py

@@ -28,6 +28,22 @@ def get_common_base(files: List[str]) -> str:
     return str(common_base)
 
 
+def check_def_node_is_nested(node) -> bool:
+    """Check if the given node is a nested function.
+
+    Args:
+        node: The node to check
+
+    Returns:
+        bool: True if the node is a nested function, False otherwise.
+    """
+    try:
+        is_nested = True if node.parent.type == "def" else False
+    except AttributeError:
+        is_nested = False
+    return is_nested
+
+
 def check_def_node_is_class_method(node) -> bool:
     """Check if the given node is a class method.