Skip to content

Commit

Permalink
feat: enables codegen create -d (#135)
Browse files Browse the repository at this point in the history
# Motivation
<!-- Why is this change necessary? -->

# Content
<!-- Please include a summary of the change -->
# Testing
<!-- How was the change tested? -->
# Please check the following before marking your PR as ready for review

- [ ] I have added tests for my changes
- [ ] I have updated the documentation or added new documentation as
needed
- [ ] I have read and agree to the [Contributor License
Agreement](../CLA.md)

---------

Co-authored-by: codegen-bot <[email protected]>
  • Loading branch information
jayhack and codegen-bot authored Jan 28, 2025
1 parent 0b8ee96 commit e4304cd
Show file tree
Hide file tree
Showing 8 changed files with 54 additions and 39 deletions.
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,9 @@

<div align="center">

[![PyPI](https://img.shields.io/pypi/v/codegen?style=flat-square&color=blue)](https://pypi.org/project/codegen/)
[![Documentation](https://img.shields.io/badge/docs-docs.codegen.com-purple?style=flat-square)](https://docs.codegen.com)
[![Slack Community](https://img.shields.io/badge/slack-community-4A154B?logo=slack&style=flat-square)](https://community.codegen.com)
[![PyPI](https://img.shields.io/pypi/v/codegen?style=flat-square&color=blue)](https://pypi.org/project/codegen/)
[![Documentation](https://img.shields.io/badge/docs-docs.codegen.com-purple?style=flat-square)](https://docs.codegen.com)
[![Slack Community](https://img.shields.io/badge/slack-community-4A154B?logo=slack&style=flat-square)](https://community.codegen.com)
[![Follow on X](https://img.shields.io/twitter/follow/codegen?style=social)](https://x.com/codegen)

</div>
Expand Down
15 changes: 12 additions & 3 deletions docs/cli/create.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@ When you run `codegen create rename-function`, it creates:
.codegen/
└── codemods/
└── rename_function/
├── rename_function.py # The codemod implementation
└── rename_function_prompt.md # System prompt (if --description used)
├── rename_function.py # The codemod implementation
└── rename_function-system-prompt.txt # System prompt (if --description used)
```

The generated codemod will have this structure:
Expand All @@ -47,7 +47,16 @@ from codegen import Codebase
def run(codebase: Codebase):
"""Add a description of what this codemod does."""
# Add your code here
pass
print('Total files: ', len(codebase.files))
print('Total functions: ', len(codebase.functions))
print('Total imports: ', len(codebase.imports))

if __name__ == "__main__":
print('Parsing codebase...')
codebase = Codebase("./")

print('Running...')
run(codebase)
```

## Examples
Expand Down
2 changes: 1 addition & 1 deletion docs/cli/notebook.mdx
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
---
title: "notebook"
title: "Notebook Command"
sidebarTitle: "notebook"
description: "Open a Jupyter notebook with the current codebase loaded"
icon: "book"
Expand Down
2 changes: 1 addition & 1 deletion src/codegen/cli/api/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ def create(self, name: str, query: str) -> CreateResponse:
return self._make_request(
"GET",
CREATE_ENDPOINT,
CreateInput(input=CreateInput.BaseCreateInput(name=name, query=query, repo_full_name=session.repo_name)),
CreateInput(input=CreateInput.BaseCreateInput(name=name, query=query, language=session.language)),
CreateResponse,
)

Expand Down
7 changes: 3 additions & 4 deletions src/codegen/cli/api/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,8 +92,8 @@ class DocsResponse(SafeBaseModel):
class CreateInput(SafeBaseModel):
class BaseCreateInput(SafeBaseModel):
name: str
query: str | None = None
repo_full_name: str | None = None
query: str
language: ProgrammingLanguage

input: BaseCreateInput

Expand All @@ -102,8 +102,7 @@ class CreateResponse(SafeBaseModel):
success: bool
response: str
code: str
codemod_id: int
context: str | None = None
context: str


###########################################################################
Expand Down
8 changes: 8 additions & 0 deletions src/codegen/cli/auth/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from codegen.cli.errors import AuthError, NoTokenError
from codegen.cli.git.repo import get_git_repo
from codegen.cli.utils.config import Config, get_config, write_config
from codegen.sdk.enums import ProgrammingLanguage


@dataclass
Expand Down Expand Up @@ -109,6 +110,13 @@ def repo_name(self) -> str:
"""Get the current repository name"""
return self.config.repo_full_name

@property
def language(self) -> ProgrammingLanguage:
"""Get the current language"""
# TODO(jayhack): This is a temporary solution to get the language.
# We should eventually get the language on init.
return self.config.programming_language or ProgrammingLanguage.PYTHON

@property
def codegen_dir(self) -> Path:
"""Get the path to the codegen-sh directory"""
Expand Down
20 changes: 11 additions & 9 deletions src/codegen/cli/codemod/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,19 +2,21 @@


def convert_to_cli(input: str, language: str, name: str) -> str:
codebase_type = "PyCodebaseType" if language.lower() == "python" else "TSCodebaseType"
return f"""import codegen.cli.sdk.decorator
# from app.codemod.compilation.models.context import CodemodContext
#from app.codemod.compilation.models.pr_options import PROptions
return f"""import codegen
from codegen import Codebase
from codegen.sdk import {codebase_type}
context: Any
@codegen.function('{name}')
def run(codebase: Codebase):
{indent(input, " ")}
@codegen.cli.sdk.decorator.function('{name}')
def run(codebase: {codebase_type}, pr_options: Any):
{indent(input, " ")}
if __name__ == "__main__":
print('Parsing codebase...')
codebase = Codebase("./")
print('Running function...')
codegen.run(run)
"""


Expand Down
33 changes: 15 additions & 18 deletions src/codegen/cli/commands/create/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
from codegen.cli.rich.codeblocks import format_command, format_path
from codegen.cli.rich.pretty_print import pretty_print_error
from codegen.cli.rich.spinners import create_spinner
from codegen.cli.utils.constants import ProgrammingLanguage
from codegen.cli.utils.default_code import DEFAULT_CODEMOD
from codegen.cli.workspace.decorators import requires_init

Expand All @@ -29,7 +28,7 @@ def get_prompts_dir() -> Path:
return PROMPTS_DIR


def get_target_path(name: str, path: Path) -> Path:
def get_target_paths(name: str, path: Path) -> tuple[Path, Path]:
"""Get the target path for the new function file.
Creates a directory structure like:
Expand All @@ -47,7 +46,9 @@ def get_target_path(name: str, path: Path) -> Path:
# Create path within .codegen/codemods
codemods_dir = base_dir / ".codegen" / "codemods"
function_dir = codemods_dir / name_snake
return function_dir / f"{name_snake}.py"
codemod_path = function_dir / f"{name_snake}.py"
prompt_path = function_dir / f"{name_snake}-system-prompt.txt"
return codemod_path, prompt_path


def make_relative(path: Path) -> str:
Expand Down Expand Up @@ -76,11 +77,11 @@ def create_command(session: CodegenSession, name: str, path: Path, description:
PATH is where to create the function (default: current directory)
"""
# Get the target path for the function
target_path = get_target_path(name, path)
codemod_path, prompt_path = get_target_paths(name, path)

# Check if file exists
if target_path.exists() and not overwrite:
rel_path = make_relative(target_path)
if codemod_path.exists() and not overwrite:
rel_path = make_relative(codemod_path)
pretty_print_error(f"File already exists at {format_path(rel_path)}\n\nTo overwrite the file:\n{format_command(f'codegen create {name} {rel_path} --overwrite')}")
return

Expand All @@ -89,34 +90,30 @@ def create_command(session: CodegenSession, name: str, path: Path, description:
try:
if description:
# Use API to generate implementation
with create_spinner("Generating function (using LLM, this will take ~30s)") as status:
with create_spinner("Generating function (using LLM, this will take ~10s)") as status:
response = RestAPI(session.token).create(name=name, query=description)
code = convert_to_cli(response.code, session.config.programming_language or ProgrammingLanguage.PYTHON, name)

# Write the system prompt if provided
if response.context:
prompt_path = get_prompts_dir() / f"{name.lower().replace(' ', '-')}-system-prompt.md"
prompt_path.write_text(response.context)
code = convert_to_cli(response.code, session.language, name)
else:
# Use default implementation
code = DEFAULT_CODEMOD.format(name=name)

# Create the target directory if needed
target_path.parent.mkdir(parents=True, exist_ok=True)
codemod_path.parent.mkdir(parents=True, exist_ok=True)

# Write the function code
target_path.write_text(code)
codemod_path.write_text(code)
prompt_path.write_text(response.context)

except (ServerError, ValueError) as e:
raise click.ClickException(str(e))

# Success message
rich.print(f"\n{'Overwrote' if overwrite and target_path.exists() else 'Created'} function '{name}'")
rich.print(f"\n{'Overwrote' if overwrite and codemod_path.exists() else 'Created'} function '{name}'")
rich.print("")
rich.print("📁 Files Created:")
rich.print(f" [dim]Function:[/dim] {make_relative(target_path)}")
rich.print(f" [dim]Function:[/dim] {make_relative(codemod_path)}")
if description and response.context:
rich.print(f" [dim]Prompt:[/dim] {make_relative(get_prompts_dir() / f'{name.lower().replace(" ", "-")}-system-prompt.md')}")
rich.print(f" [dim]Prompt:[/dim] {make_relative(prompt_path)}")

# Next steps
rich.print("\n[bold]What's next?[/bold]\n")
Expand Down

0 comments on commit e4304cd

Please sign in to comment.