🎨 black

.gitignore

@@ -7,3 +7,4 @@ site/
 test.py
 coverage.xml
 .DS_Store
+wizard/

tests/data/calculate.py

@@ -1,5 +1,6 @@
 class Calculate:
     def multiply(a, b):
         return a * b
+
     def add(a, b):
-        return a * b
+        return a * b

tests/data/multiply.py

@@ -1,2 +1,2 @@
 def multiply(a, b):
-    return a * b
+    return a * b

tests/data/multiply_docstring.py

@@ -10,4 +10,4 @@ def multiply(a, b):
       >>> multiply(2, 3)
       6
     """
-    return a * b
+    return a * b

tests/test_cli_main.py

@@ -7,7 +7,7 @@
 import unittest.mock as mock
 
 
-@pytest.fixture(scope='function')
+@pytest.fixture(scope="function")
 def file_path(tmp_path) -> Path:
     temp_file = tmp_path / "test_add.py"
     temp_file.write_text("def add(a, b):\n  return a + b")
@@ -36,7 +36,10 @@ def test_callback_version():
         (False, True, True, True),
     ],
 )
-@mock.patch('write_the.llm.LLM.run', return_value="\n\nadd:\n  Sums 2 numbers.\n  Args:\n    a (int): The first number to add.\n    b (int): The second number to add.\n  Returns:\n    int: The sum of `a` and `b`.\n  Examples:\n    >>> add(1, 2)\n    3\n\n")
+@mock.patch(
+    "write_the.llm.LLM.run",
+    return_value="\n\nadd:\n  Sums 2 numbers.\n  Args:\n    a (int): The first number to add.\n    b (int): The second number to add.\n  Returns:\n    int: The sum of `a` and `b`.\n  Examples:\n    >>> add(1, 2)\n    3\n\n",
+)
 def test_docs_mocked(mocked_run, file_path: Path, nodes, save, context, pretty, force):
     runner = CliRunner()
     args = ["docs", str(file_path)]
@@ -75,9 +78,9 @@ def test_mkdocs(tmp_path: Path):
     print(result.stdout)
     assert result.exit_code == 0
     files = [f.name for f in tmp_path.glob("*")]
-    assert 'mkdocs.yml' in files
-    assert '.github' in files
-    assert 'docs' in files
+    assert "mkdocs.yml" in files
+    assert ".github" in files
+    assert "docs" in files
 
 
 @pytest.mark.parametrize(
@@ -90,14 +93,17 @@ def test_mkdocs(tmp_path: Path):
         (False, True, True),
     ],
 )
-@mock.patch('write_the.llm.LLM.run', return_value="""@pytest.mark.parametrize(
+@mock.patch(
+    "write_the.llm.LLM.run",
+    return_value="""@pytest.mark.parametrize(
     "a, b, expected", [(2, 3, 5), (0, 5, 5), (-2, -3, -5), (2.5, 3, 5.5), (2, -3, -1)]
 )
 def test_add(a, b, expected):
-    assert add(a, b) == expected""")
+    assert add(a, b) == expected""",
+)
 def test_tests_mocked(mocked_run, file_path: Path, save, pretty, force):
     runner = CliRunner()
-    test_dir = file_path.parent / 'docs'
+    test_dir = file_path.parent / "docs"
     args = ["tests", str(file_path), "--out", test_dir]
 
     if save:
@@ -117,4 +123,4 @@ def test_tests_mocked(mocked_run, file_path: Path, save, pretty, force):
         assert "assert add(a, b) == expected" in test_file.read_text()
         assert str(file_path) in result.stdout
     else:
-        assert "assert add(a, b) == expected" in result.stdout
+        assert "assert add(a, b) == expected" in result.stdout

tests/test_cst_docstring_adder.py

@@ -22,28 +22,28 @@ def function_def_node():
     return cst.FunctionDef(
         name=cst.Name("function_name"),
         params=cst.Parameters(),
-        body=cst.IndentedBlock(
-            body=[cst.SimpleStatementLine(body=[cst.Pass()])]),
+        body=cst.IndentedBlock(body=[cst.SimpleStatementLine(body=[cst.Pass()])]),
     )
 
 
 @pytest.fixture
 def class_def_node():
-    return cst.ClassDef(name=cst.Name("ClassName"), body=cst.IndentedBlock(body=[cst.SimpleStatementLine(body=[cst.Pass()])]))
+    return cst.ClassDef(
+        name=cst.Name("ClassName"),
+        body=cst.IndentedBlock(body=[cst.SimpleStatementLine(body=[cst.Pass()])]),
+    )
 
 
 @pytest.fixture
 def method_def_node():
     method_def = cst.FunctionDef(
         name=cst.Name("method_name"),
         params=cst.Parameters(params=[cst.Param(name=cst.Name("cls"))]),
-        body=cst.IndentedBlock(
-            body=[cst.SimpleStatementLine(body=[cst.Pass()])]),
+        body=cst.IndentedBlock(body=[cst.SimpleStatementLine(body=[cst.Pass()])]),
     )
     return cst.ClassDef(
         name=cst.Name("ClassName"),
-        body=cst.IndentedBlock(
-            body=[cst.SimpleStatementLine(body=[method_def])]),
+        body=cst.IndentedBlock(body=[cst.SimpleStatementLine(body=[method_def])]),
     )
 
 
@@ -66,16 +66,14 @@ def test_leave_function_def_without_docstring(docstrings, force, function_def_no
 
 def test_leave_class_def_with_docstring(docstrings, force, class_def_node):
     docstring_adder = DocstringAdder(docstrings, force)
-    updated_node = docstring_adder.leave_ClassDef(
-        class_def_node, class_def_node)
+    updated_node = docstring_adder.leave_ClassDef(class_def_node, class_def_node)
     assert has_docstring(updated_node) is False
 
 
 def test_leave_class_def_without_docstring(docstrings, force, class_def_node):
     docstrings.pop("ClassName.method_name")
     docstring_adder = DocstringAdder(docstrings, force)
-    updated_node = docstring_adder.leave_ClassDef(
-        class_def_node, class_def_node)
+    updated_node = docstring_adder.leave_ClassDef(class_def_node, class_def_node)
     assert not has_docstring(updated_node)
 
 
@@ -84,10 +82,8 @@ def test_leave_method_def_without_docstring(
 ):
     docstrings.pop("ClassName.method_name")
     docstring_adder = DocstringAdder(docstrings, force)
-    updated_node = docstring_adder.leave_ClassDef(
-        class_def_node, class_def_node)
-    updated_node = docstring_adder.leave_FunctionDef(
-        method_def_node, method_def_node)
+    updated_node = docstring_adder.leave_ClassDef(class_def_node, class_def_node)
+    updated_node = docstring_adder.leave_FunctionDef(method_def_node, method_def_node)
     assert not has_docstring(updated_node)
 
 
@@ -110,10 +106,14 @@ def test_add_docstring_with_force(docstrings, function_def_node):
     updated_node = docstring_adder.add_docstring(function_def_node)
     assert has_docstring(updated_node)
 
+
 def test_add_docstring_escape_newline(docstrings, function_def_node):
     force = True
-    docstrings['function_name'] = """\\ntest\ntest\\\\n\\n"""
+    docstrings["function_name"] = """\\ntest\ntest\\\\n\\n"""
     docstring_adder = DocstringAdder(docstrings, force)
     updated_node = docstring_adder.add_docstring(function_def_node)
     assert has_docstring(updated_node)
-    assert get_docstring(updated_node).strip('"""').strip() == """\\\\ntest\n    test\\\\n\\\\n"""
+    assert (
+        get_docstring(updated_node).strip('"""').strip()
+        == """\\\\ntest\n    test\\\\n\\\\n"""
+    )

tests/test_cst_docstring_remover.py

@@ -3,6 +3,7 @@
 from write_the.cst.docstring_remover import DocstringRemover, remove_docstrings
 from write_the.cst.utils import get_docstring
 
+
 @pytest.fixture
 def tree():
     return cst.parse_module(
@@ -34,7 +35,7 @@ def nodes():
 def test_leave_FunctionDef(tree, nodes):
     remover = DocstringRemover(nodes)
     updated_tree = tree.visit(remover)
-    assert get_docstring(updated_tree.body[0]) is None 
+    assert get_docstring(updated_tree.body[0]) is None
     assert get_docstring(updated_tree.body[1]) == "'''This is another docstring.'''"
 
 

tests/test_cst_function_and_class_collector.py

@@ -41,8 +41,8 @@ def test_visit_ClassDef_with_docstring(tree):
 
 
 def test_get_node_names(tree, force):
-    assert get_node_names(tree, force) == ["foo", "Bar"]
+    assert get_node_names(tree, force) == ["Bar", "foo"]
 
 
 def test_get_node_names_with_force_true(tree, force):
-    assert get_node_names(tree, True) == ["foo", "Bar"]
+    assert get_node_names(tree, True) == ["Bar", "foo"]

tests/test_cst_node_extractor.py

@@ -2,6 +2,7 @@
 import libcst as cst
 from write_the.cst.node_extractor import NodeExtractor, extract_nodes_from_tree
 
+
 @pytest.fixture
 def tree():
     return cst.parse_module(

tests/test_cst_utils.py

@@ -2,31 +2,42 @@
 from write_the.cst.utils import has_docstring, nodes_to_tree
 import pytest
 
+
 @pytest.fixture
 def cst_function_def():
     return cst.FunctionDef(
-    name=cst.Name("function_name"),
-    params=cst.Parameters(),
-    body=cst.IndentedBlock(
-        body=[
-            cst.SimpleStatementLine(body=[cst.Expr(value=cst.SimpleString('"""This is a docstring."""'))]),
-            cst.SimpleStatementLine(body=[cst.Pass()]),
-        ]
-    ),
-)
+        name=cst.Name("function_name"),
+        params=cst.Parameters(),
+        body=cst.IndentedBlock(
+            body=[
+                cst.SimpleStatementLine(
+                    body=[
+                        cst.Expr(value=cst.SimpleString('"""This is a docstring."""'))
+                    ]
+                ),
+                cst.SimpleStatementLine(body=[cst.Pass()]),
+            ]
+        ),
+    )
 
 
 @pytest.fixture
 def cst_class_def():
     return cst.ClassDef(
-    name=cst.Name("ClassName"),
-    body=cst.IndentedBlock(
-        body=[
-            cst.SimpleStatementLine(body=[cst.Expr(value=cst.SimpleString('"""This is a class docstring."""'))]),
-            cst.SimpleStatementLine(body=[cst.Pass()]),
-        ]
-    ),
-)
+        name=cst.Name("ClassName"),
+        body=cst.IndentedBlock(
+            body=[
+                cst.SimpleStatementLine(
+                    body=[
+                        cst.Expr(
+                            value=cst.SimpleString('"""This is a class docstring."""')
+                        )
+                    ]
+                ),
+                cst.SimpleStatementLine(body=[cst.Pass()]),
+            ]
+        ),
+    )
 
 
 def test_has_docstring_function_def(cst_function_def):

write_the/cli/__init__.py

@@ -1 +1 @@
-from .main import app
+from .main import app

write_the/commands/__init__.py

@@ -1,3 +1,3 @@
 from .docs import write_the_docs
 from .mkdocs import write_the_mkdocs
-from .tests import write_the_tests
+from .tests import write_the_tests

write_the/commands/docs/__init__.py

@@ -1 +1 @@
-from .docs import write_the_docs
+from .docs import write_the_docs

write_the/commands/docs/prompts.py

@@ -26,4 +26,6 @@ def add(a, b):
 {code}
 Formatted docstrings for {nodes}:
 """
-write_docstings_for_nodes_prompt = PromptTemplate(input_variables=["code", "nodes"], template=docs_template)
+write_docstings_for_nodes_prompt = PromptTemplate(
+    input_variables=["code", "nodes"], template=docs_template
+)

write_the/commands/docs/utils.py

@@ -1,22 +1,24 @@
 import libcst as cst
-import re 
+import re
 
 from write_the.cst import nodes_to_tree
 from write_the.cst.function_and_class_collector import get_node_names
 from write_the.cst.node_extractor import extract_nodes_from_tree
 from write_the.cst.node_remover import remove_nodes_from_tree
 
+
 def pad_with_newline_if_needed(s):
-    if not s.startswith('\n'):
-        s = '\n' + s
-    if not s.endswith('\n'):
-        s = s + '\n'
+    if not s.startswith("\n"):
+        s = "\n" + s
+    if not s.endswith("\n"):
+        s = s + "\n"
     return s
 
+
 def extract_block(text, class_function_names):
     results = {}
     for name in class_function_names:
-        pattern = rf'({name}:[\s\S]*?)(?=(\n\w|\Z))'
+        pattern = rf"({name}:[\s\S]*?)(?=(\n\w|\Z))"
         match = re.search(pattern, text)
         if match:
             block = match.group(1).lstrip(f"{name}:")
@@ -43,14 +45,11 @@ def process_nodes(tree: cst.Module, nodes, context, extract_specific_nodes) -> s
             extracted_nodes = extract_nodes_from_tree(tree, nodes)
             processed_tree = nodes_to_tree(extracted_nodes)
         else:
-            all_nodes = get_node_names(tree, False)
+            all_nodes = get_node_names(tree, True)
             nodes_to_remove = [n for n in all_nodes if n not in nodes]
-            processed_tree = remove_nodes_from_tree(
-                tree, nodes_to_remove
-            )
+            processed_tree = remove_nodes_from_tree(tree, nodes_to_remove)
         code = processed_tree.code
     else:
         code = tree.code
 
     return code
-

write_the/commands/mkdocs/mkdocs.py

@@ -4,7 +4,10 @@
 
 from .templates import action_template, mkdocs_template
 
-def write_the_mkdocs(code_dir: Path, readme: Path = None, out_dir: Path = Path('.'), project_name=None):
+
+def write_the_mkdocs(
+    code_dir: Path, readme: Path = None, out_dir: Path = Path("."), project_name=None
+):
     """
     Generates a mkdocs project from a directory of python files.
     Args:
@@ -33,12 +36,14 @@ def write_the_mkdocs(code_dir: Path, readme: Path = None, out_dir: Path = Path('
             continue
         key = "index"
         for group in groups:
-            if f"{code_dir.name}/{group}/" in str(file) or f"{code_dir.name}/{group}." in str(file):
+            if f"{code_dir.name}/{group}/" in str(
+                file
+            ) or f"{code_dir.name}/{group}." in str(file):
                 key = group
                 break
         module = str(file).rstrip(".py").replace("/", ".")  # breaks on windows?
         references[key].append(f"::: {module}")
-    docs_dir = out_dir / 'docs'
+    docs_dir = out_dir / "docs"
     reference_path = docs_dir / "reference"
     reference_path.mkdir(parents=True, exist_ok=True)
     for doc in references:

write_the/commands/mkdocs/templates.py

@@ -54,4 +54,4 @@
           path: .cache
       - run: pip install "mkdocstrings[python]" "mkdocs-material"
       - run: mkdocs gh-deploy --force
-"""
+"""

write_the/commands/tests/__init__.py

@@ -1 +1 @@
-from .tests import write_the_tests
+from .tests import write_the_tests

write_the/commands/tests/prompts.py

@@ -12,6 +12,6 @@
 {code}
 ```
 """
-write_tests_for_file_prompt = PromptTemplate(input_variables=["code", "path"], template=tests_template)
-
-
+write_tests_for_file_prompt = PromptTemplate(
+    input_variables=["code", "path"], template=tests_template
+)