Commands

write_the_tests(filename, model='gpt-3.5-turbo-instruct') async

Formats and runs the tests for a given file using a specified model.

Parameters:

Name	Type	Description	Default
filename	Path	The path to the file to be tested.	required
model	str	The model to use for the generation. Defaults to "gpt-3.5-turbo-instruct".	'gpt-3.5-turbo-instruct'

Returns:

Name	Type	Description
str	str	The formatted and tested code.

Examples:

>>> write_the_tests(Path("test.py"), "gpt-3.5-turbo-instruct")
"Formatted and tested code"

Note

This function is asynchronous and should be awaited when called.

Source code in write_the/commands/tests/tests.py

async def write_the_tests(filename: Path, model="gpt-3.5-turbo-instruct") -> str:
    """
    Formats and runs the tests for a given file using a specified model.

    Args:
      filename (Path): The path to the file to be tested.
      model (str): The model to use for the generation. Defaults to "gpt-3.5-turbo-instruct".

    Returns:
      str: The formatted and tested code.

    Examples:
      >>> write_the_tests(Path("test.py"), "gpt-3.5-turbo-instruct")
      "Formatted and tested code"

    Note:
      This function is asynchronous and should be awaited when called.
    """
    with open(filename, "r") as file:
        source_code = file.read()
    source_code = format_str(source_code, mode=FileMode())
    llm = LLM(write_tests_for_file_prompt, model_name=model)
    result = await llm.run(code=source_code, path=filename)
    code = (
        result.strip()
        .lstrip("Test Code:\n```python")
        .lstrip("```python")
        .lstrip("```")
        .rstrip("```")
    )
    return format_str(code, mode=FileMode())

write_the_converters(filename, input_format, output_format, model='gpt-3.5-turbo-instruct') async

Formats and runs the tests for a given file.

Parameters:

Name	Type	Description	Default
filename	Path	The path to the file to be tested.	required
input_format	str	The input format of the file.	required
output_format	str	The format to convert the file to.	required
model	str	The model to use for conversion. Defaults to "gpt-3.5-turbo-instruct".	'gpt-3.5-turbo-instruct'

Returns:

Name	Type	Description
str	str	The converted output.

Examples:

>>> write_the_converters(Path(".travis.yml"), input_format="Travis CI", output_format="Github Actions", model="gpt-3.5-turbo-instruct")
"The converted output"

Source code in write_the/commands/converters/converters.py

async def write_the_converters(filename: Path, input_format: str, output_format: str, model: str = "gpt-3.5-turbo-instruct") -> str:
    """
    Formats and runs the tests for a given file.

    Args:
      filename (Path): The path to the file to be tested.
      input_format (str): The input format of the file.
      output_format (str): The format to convert the file to.
      model (str, optional): The model to use for conversion. Defaults to "gpt-3.5-turbo-instruct".

    Returns:
      str: The converted output.

    Examples:
      >>> write_the_converters(Path(".travis.yml"), input_format="Travis CI", output_format="Github Actions", model="gpt-3.5-turbo-instruct")
      "The converted output"
    """
    with open(filename, "r") as file:
        source_text = file.read()

    llm = LLM(write_converters_for_file_prompt, model_name=model)
    result = await llm.run(code=source_text, input_format=input_format, output_format=output_format)

    formatted_text = result.strip()

    if formatted_text.endswith('```') and not source_text.endswith('```'):
        # strip the last line of the code block if the source text didn't end with a code block
        formatted_text = formatted_text[: formatted_text.rfind("\n")]
    if formatted_text.startswith('```') and not source_text.startswith('```'):
        # strip the first line of the code block if the source text didn't start with a code block
        formatted_text = formatted_text[formatted_text.find("\n") + 1 :]
    return formatted_text.strip()

write_the_docs(tree, node_names=[], update=False, force=False, save=False, context=False, background=True, pretty=False, max_batch_size=False, model='gpt-3.5-turbo-instruct') async

Generates docstrings for a given tree of nodes using a specified model.

Parameters:

Name	Type	Description	Default
tree	cst.Module	The tree of nodes to write docs for.	required
node_names	list	The list of nodes names to write docs for. Defaults to an empty list.	[]
update	bool	Whether to update existing docstrings. Defaults to False.	False
force	bool	Whether to force writing of docs. Defaults to False.	False
save	bool	Whether to save the docs. Defaults to False.	False
context	bool	Whether to include context nodes. Defaults to False.	False
background	bool	Whether to run the process in the background. Defaults to True.	True
pretty	bool	Whether to format the code. Defaults to False.	False
max_batch_size	bool	Max number of nodes in each batch. Defaults to False.	False
model	str	The model to use for the generation. Defaults to "gpt-3.5-turbo-instruct".	'gpt-3.5-turbo-instruct'

Returns:

Name	Type	Description
str	str	The source code with the generated docstrings.

Raises:

Type	Description
FileSkippedError	If no nodes are found.

Notes

If node_names is provided, force is set to True and context is set to False.

Examples:

>>> write_the_docs(tree, model="gpt-3.5-turbo-instruct")
"def add(a, b):
    """Sums 2 numbers.
    Args:
        a (int): The first number to add.
        b (int): The second number to add.
    Returns:
        int: The sum of `a` and `b`.
    """
    return a + b"

Source code in write_the/commands/docs/docs.py

async def write_the_docs(
    tree: cst.Module,
    node_names=[],
    update=False,
    force=False,
    save=False,
    context=False,
    background=True,
    pretty=False,
    max_batch_size=False,
    model="gpt-3.5-turbo-instruct",
) -> str:
    """
    Generates docstrings for a given tree of nodes using a specified model.

    Args:
      tree (cst.Module): The tree of nodes to write docs for.
      node_names (list, optional): The list of nodes names to write docs for. Defaults to an empty list.
      update (bool, optional): Whether to update existing docstrings. Defaults to False.
      force (bool, optional): Whether to force writing of docs. Defaults to False.
      save (bool, optional): Whether to save the docs. Defaults to False.
      context (bool, optional): Whether to include context nodes. Defaults to False.
      background (bool, optional): Whether to run the process in the background. Defaults to True.
      pretty (bool, optional): Whether to format the code. Defaults to False.
      max_batch_size (bool, optional): Max number of nodes in each batch. Defaults to False.
      model (str, optional): The model to use for the generation. Defaults to "gpt-3.5-turbo-instruct".

    Returns:
      str: The source code with the generated docstrings.

    Raises:
      FileSkippedError: If no nodes are found.

    Notes:
      If `node_names` is provided, `force` is set to `True` and `context` is set to `False`.

    Examples:
      >>> write_the_docs(tree, model="gpt-3.5-turbo-instruct")
      "def add(a, b):
          \"\"\"Sums 2 numbers.
          Args:
              a (int): The first number to add.
              b (int): The second number to add.
          Returns:
              int: The sum of `a` and `b`.
          \"\"\"
          return a + b"
    """
    extract_specific_nodes = False
    if node_names:
        extract_specific_nodes = True
        force = True
    else:
        node_names = get_node_names(tree, force=force, update=update)
    if not node_names:
        raise FileSkippedError("No nodes found, skipping file...")
    if update:
        remove_docstrings = False
        llm = LLM(update_docstrings_for_nodes_prompt, model_name=model)
    else:
        remove_docstrings = True
        llm = LLM(write_docstrings_for_nodes_prompt, model_name=model)

    batches = create_batches(
        tree=tree,
        node_names=node_names,
        max_tokens=llm.max_tokens,
        prompt_size=llm.prompt_size,
        response_size_per_node=250,  # a guess... TODO: smarter
        max_batch_size=max_batch_size,
        send_background_context=background,
        send_node_context=context,
        remove_docstrings=remove_docstrings
    )
    promises = []
    node_names_list = []
    for batch in batches:
        node_names = batch.node_names
        code = batch.code
        promises.append((llm.run(code=code, nodes=node_names)))
        node_names_list.append(node_names)
    # Can i yield here so batches can be logged?
    results = await asyncio.gather(*promises)
    docstring_dict = {}
    for node_names, result in zip(node_names_list, results):
        docstring_dict.update(extract_block(result, node_names))
    modified_tree = add_docstrings_to_tree(tree, docstring_dict, force=force or update)
    if not save and extract_specific_nodes:
        extracted_nodes = extract_nodes_from_tree(modified_tree, node_names)
        modified_tree = nodes_to_tree(extracted_nodes)
    if pretty:
        return format_str(modified_tree.code, mode=FileMode())
    return modified_tree.code

process_nodes(tree, nodes, context, extract_specific_nodes)

Processes a tree of nodes.

Parameters:

Name	Type	Description	Default
tree	cst.Module	The tree of nodes to process.	required
nodes	list	The list of nodes to process.	required
context	bool	Whether to include context nodes.	required
extract_specific_nodes	bool	Whether to extract specific nodes.	required

Returns:

Name	Type	Description
str	str	The processed tree as a string.

Examples:

>>> process_nodes(tree, nodes, context, extract_specific_nodes)
"Processed tree as a string"

Source code in write_the/commands/docs/utils.py

def process_nodes(tree: cst.Module, nodes, context, extract_specific_nodes) -> str:
    """
    Processes a tree of nodes.

    Args:
      tree (cst.Module): The tree of nodes to process.
      nodes (list): The list of nodes to process.
      context (bool): Whether to include context nodes.
      extract_specific_nodes (bool): Whether to extract specific nodes.

    Returns:
      str: The processed tree as a string.

    Examples:
      >>> process_nodes(tree, nodes, context, extract_specific_nodes)
      "Processed tree as a string"
    """
    if not context:
        if extract_specific_nodes:
            extracted_nodes = extract_nodes_from_tree(tree, nodes)
            processed_tree = nodes_to_tree(extracted_nodes)
        else:
            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)
        code = processed_tree.code
    else:
        code = tree.code

    return code

write_the_mkdocs(code_dir, readme=None, out_dir=Path('.'), project_name=None)

Generates a mkdocs project from a directory of python files.

Parameters:

Name	Type	Description	Default
code_dir	Path	The directory containing the python files.	required
readme	Path	The readme file to include in the project. Defaults to None.	None
out_dir	Path	The directory to write the project to. Defaults to the current directory.	Path('.')
project_name	str	The name of the project. Defaults to the name of the code_dir.	None

Notes

If readme is not provided, the project will not have a home page. If project_name is not provided, the project will be named after the code_dir.

Side Effects

Creates a mkdocs project in the out_dir. Creates a .github/workflows/mkdocs.yml file in the out_dir.

Returns:

Type	Description
	None

Source code in write_the/commands/mkdocs/mkdocs.py

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:
      code_dir (Path): The directory containing the python files.
      readme (Path, optional): The readme file to include in the project. Defaults to None.
      out_dir (Path, optional): The directory to write the project to. Defaults to the current directory.
      project_name (str, optional): The name of the project. Defaults to the name of the code_dir.

    Notes:
      If readme is not provided, the project will not have a home page.
      If project_name is not provided, the project will be named after the code_dir.

    Side Effects:
      Creates a mkdocs project in the out_dir.
      Creates a .github/workflows/mkdocs.yml file in the out_dir.

    Returns:
      None
    """
    files = list_python_files(code_dir)
    groups = [path.stem for path in code_dir.glob("*") if not path.stem.startswith("_")]

    if not project_name:
        project_name = code_dir.name
    mkdocs = mkdocs_template.format(project_name=project_name)
    references = defaultdict(list)
    for file in files:
        if file.name.startswith("_"):
            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):
                key = group
                break

        module = str(file).rsplit(".py")[0]
        if sys.platform.startswith("win32") or sys.platform.startswith("cygwin"):
            splitter = "\\"
        else:
            splitter = "/"
        module = module.replace(splitter, ".")

        references[key].append(f"::: {module}")
    docs_dir = out_dir / "docs"
    reference_path = docs_dir / "reference"
    reference_path.mkdir(parents=True, exist_ok=True)
    for doc in references:
        with open(f"{reference_path}/{doc}.md", "w") as f:
            for ref in references[doc]:
                f.write(ref + "\n\n")
    if readme:
        index_text = f"---\ntitle: Home\n---\n{readme.read_text()}"
        (docs_dir / "index.md").write_text(index_text)
    if not (out_dir / "mkdocs.yml").exists():
        (out_dir / "mkdocs.yml").write_text(mkdocs)
    action_path = out_dir / ".github" / "workflows" / "mkdocs.yml"
    if not action_path.exists():
        action_path.parent.mkdir(parents=True, exist_ok=True)
        action_path.write_text(action_template)