Merge pull request #43 from MichaelisTrofficus/new-traversing-structure

New traversing structure

Author: MichaelisTrofficus

.flake8

@@ -1,6 +1,6 @@
 [flake8]
 select = B,B9,C,D,DAR,E,F,N,RST,S,W
-ignore = E203,E501,RST201,RST203,RST301,W503,D100,D103,D102,D101,D107,D105,D205,D212,D415,D104,N803,S605
+ignore = E203,E501,RST201,RST203,RST301,W503,D100,D103,D102,D101,D107,D105,D205,D212,D415,D104,N803,S605,N802,C901,B904
 max-line-length = 120
 max-complexity = 10
 docstring-convention = google

.gitignore

@@ -328,3 +328,4 @@ pyrightconfig.json
 chatgpt_experiment.py
 chatgpt_battery.py
 chatgpt_battery_classes.py
+*.diff

noxfile.py

@@ -159,7 +159,7 @@ def coverage(session: Session) -> None:
         "coverage",
         *args,
         "--omit=tests/*,src/gpt4docstrings/__main__.py,src/gpt4docstrings/cli.py",
-        "--fail-under=40",
+        "--fail-under=30",
     )
 
 

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "gpt4docstrings"
-version = "0.1.2"
+version = "0.1.3"
 description = "gpt4docstrings"
 authors = ["Miguel Otero Pedrido <miguel.otero.pedrido.1993@gmail.com>"]
 license = "MIT"
@@ -25,6 +25,11 @@ pytest-mock = "^3.11.1"
 tabulate = "^0.9.0"
 colorama = "^0.4.6"
 langchain = "^0.0.311"
+toml = "^0.10.2"
+asyncio = "^3.4.3"
+aiofiles = "^23.2.1"
+reorder-python-imports = "^3.12.0"
+astor = "^0.8.1"
 
 [tool.poetry.dev-dependencies]
 Pygments = ">=2.10.0"

pyproject.toml

@@ -3,6 +3,7 @@
 import click
 
 import gpt4docstrings
+from gpt4docstrings.config import GPT4DocstringsConfig
 
 
 @click.option(
@@ -73,5 +74,6 @@ def main(paths, **kwargs):
         docstring_style=kwargs["style"],
         api_key=kwargs["api_key"],
         verbose=kwargs["verbose"],
+        config=GPT4DocstringsConfig(),
     )
     docstrings_generator.generate_docstrings()

src/gpt4docstrings/cli.py

@@ -0,0 +1,66 @@
+# Copyright 2020 Lynn Root
+"""Configuration-related helpers."""
+# Adapted from Black https://github.com/psf/black/blob/master/black.py.
+# This code is adapted from Lynn Root's interrogate library: https://github.com/econchick/interrogate
+import pathlib
+
+import attr
+
+
+# TODO: idea: break out InterrogateConfig into two classes: one for
+# running the tool, one for reporting the results
+@attr.s
+class GPT4DocstringsConfig:
+    """
+    Configuration related to interrogating a given codebase.
+
+    Args:
+        update_file (bool): If `True`, the documented file will be updated in place. This may be dangerous
+            in same cases, so it's better to use patches to confirm the changes to be applied.
+        ignore_private (bool): Ignore private classes, methods, and
+            functions starting with two underscores.
+        ignore_semiprivate (bool): Ignore semiprivate classes, methods,
+            and functions starting with a single underscore.
+        ignore_init_method (bool): Ignore ``__init__`` method of
+            classes.
+        ignore_nested_functions (bool): Ignore nested functions and
+            methods.
+    """
+
+    update_file = attr.ib(default=False)
+    ignore_private = attr.ib(default=False)
+    ignore_semiprivate = attr.ib(default=False)
+    ignore_init_method = attr.ib(default=False)
+    ignore_nested_classes = attr.ib(default=False)
+    ignore_nested_functions = attr.ib(default=False)
+    ignore_property_setters = attr.ib(default=False)
+    ignore_property_decorators = attr.ib(default=False)
+
+
+def find_project_root(srcs):
+    """
+    Return a directory containing .git, .hg, or pyproject.toml.
+    That directory can be one of the directories passed in `srcs` or their
+    common parent.
+    If no directory in the tree contains a marker that would specify it's the
+    project root, the root of the file system is returned.k
+    """
+    if not srcs:
+        return pathlib.Path("/").resolve()
+
+    common_base = min(pathlib.Path(src).resolve() for src in srcs)
+    if common_base.is_dir():
+        # Append a fake file so `parents` below returns `common_base_dir`, too.
+        common_base /= "fake-file"
+
+    for directory in common_base.parents:
+        if (directory / ".git").exists():
+            return directory
+
+        if (directory / ".hg").is_dir():
+            return directory
+
+        if (directory / "pyproject.toml").is_file():
+            return directory
+
+    return directory

src/gpt4docstrings/config.py

@@ -1,17 +1,16 @@
+import ast
 import os
 import textwrap
 
 import openai
 from langchain.chat_models import ChatOpenAI
 from langchain.prompts import PromptTemplate
-from redbaron import RedBaron
 
-from gpt4docstrings import utils
-from gpt4docstrings.docstrings_generators.utils.decorators import retry
+from gpt4docstrings.docstrings_generators.docstring import Docstring
 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
+from gpt4docstrings.visit import GPT4DocstringsNode
 
 
 class ChatGPTDocstringGenerator:
@@ -32,12 +31,13 @@ def __init__(
         self.model_name = model_name
         self.docstring_style = docstring_style
 
-        self.model = ChatOpenAI(model_name=model_name, temperature=1.0)
+        self.model = ChatOpenAI(
+            model_name=model_name, temperature=1.0, openai_api_key=self.api_key
+        )
         self.function_prompt_template = FUNCTION_PROMPTS.get(docstring_style)
         self.class_prompt_template = CLASS_PROMPTS.get(docstring_style)
 
-    @retry(max_retries=5, delay=5)
-    def _get_completion(self, prompt: str) -> str:
+    async def _get_completion(self, prompt: str) -> str:
         """
         Generates a completion using the ChatGPT model.
 
@@ -47,86 +47,42 @@ def _get_completion(self, prompt: str) -> str:
         Returns:
             str: The generated completion.
         """
-        return self.model.predict(prompt).strip()
+        return await self.model.apredict(prompt)
+
+    def _get_template(self, node: GPT4DocstringsNode):
+        """Returns a function template or a class template depending on the node type"""
+        if node.node_type in ["FunctionDef", "AsyncFunctionDef"]:
+            return self.function_prompt_template
+        else:
+            return self.class_prompt_template
 
-    def generate_function_docstring(self, source: str) -> dict:
+    async def generate_docstring(self, node: GPT4DocstringsNode) -> Docstring:
         """
         Generates a docstring for a function.
 
         Args:
-            source (str): The source code of the function.
+            node (GPT4DocstringsNode): A GPT4DocstringsNode node
 
         Returns:
-            dict: A dictionary containing the generated docstring.
-
-        Raises:
-            ASTError: Raises an ASTError when there are errors interacting with an AST node
+            Docstring: A Docstring object
         """
-        source = source.strip()
+        source = node.source.strip()
         stripped_source = textwrap.dedent(source)
-        prompt = PromptTemplate(
-            template=self.function_prompt_template,
-            input_variables=["code"],
-        )
-        _input = prompt.format_prompt(code=stripped_source)
-        fn_src = DocstringParser().parse(self._get_completion(_input.to_string()))
-
-        try:
-            fn_node = RedBaron(fn_src).find_all("def")[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:
-        """
-        Generates docstrings for a class.
+        prompt_template = self._get_template(node)
+        parent_offset = node.col_offset
 
-        Args:
-            source (str): The source code of the class.
-
-        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)
         prompt = PromptTemplate(
-            template=self.class_prompt_template,
+            template=prompt_template,
             input_variables=["code"],
         )
         _input = prompt.format_prompt(code=stripped_source)
-        class_src = DocstringParser().parse(self._get_completion(_input.to_string()))
-
-        # TODO: Add here access to class node explicitly.
-        try:
-            class_node = RedBaron(class_src).find_all("class")[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,
+        src = DocstringParser().parse(await self._get_completion(_input.to_string()))
+
+        tree = ast.parse(src)
+        for n in ast.walk(tree):
+            if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                return Docstring(
+                    text=ast.get_docstring(n),
+                    col_offset=n.body[-1].col_offset + parent_offset,
+                    lineno=node.lineno,
                 )
-
-            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
-
-        except ValueError as e:
-            raise ASTError(
-                f"Some error has occurred when trying to parse the current AST node: {e}"
-            ) from e