Skip to content

Commit 03356da

Browse files
jayhackcodegen-bot
and
codegen-bot
authored
feat: enables codegen create -d (#135)
# 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]>
1 parent 90484ec commit 03356da

File tree

8 files changed

+54
-39
lines changed

8 files changed

+54
-39
lines changed

README.md

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,9 @@
1616

1717
<div align="center">
1818

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

2424
</div>

docs/cli/create.mdx

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,8 @@ When you run `codegen create rename-function`, it creates:
3333
.codegen/
3434
└── codemods/
3535
└── rename_function/
36-
├── rename_function.py # The codemod implementation
37-
└── rename_function_prompt.md # System prompt (if --description used)
36+
├── rename_function.py # The codemod implementation
37+
└── rename_function-system-prompt.txt # System prompt (if --description used)
3838
```
3939

4040
The generated codemod will have this structure:
@@ -47,7 +47,16 @@ from codegen import Codebase
4747
def run(codebase: Codebase):
4848
"""Add a description of what this codemod does."""
4949
# Add your code here
50-
pass
50+
print('Total files: ', len(codebase.files))
51+
print('Total functions: ', len(codebase.functions))
52+
print('Total imports: ', len(codebase.imports))
53+
54+
if __name__ == "__main__":
55+
print('Parsing codebase...')
56+
codebase = Codebase("./")
57+
58+
print('Running...')
59+
run(codebase)
5160
```
5261

5362
## Examples

docs/cli/notebook.mdx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
---
2-
title: "notebook"
2+
title: "Notebook Command"
33
sidebarTitle: "notebook"
44
description: "Open a Jupyter notebook with the current codebase loaded"
55
icon: "book"

src/codegen/cli/api/client.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ def create(self, name: str, query: str) -> CreateResponse:
174174
return self._make_request(
175175
"GET",
176176
CREATE_ENDPOINT,
177-
CreateInput(input=CreateInput.BaseCreateInput(name=name, query=query, repo_full_name=session.repo_name)),
177+
CreateInput(input=CreateInput.BaseCreateInput(name=name, query=query, language=session.language)),
178178
CreateResponse,
179179
)
180180

src/codegen/cli/api/schemas.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -92,8 +92,8 @@ class DocsResponse(SafeBaseModel):
9292
class CreateInput(SafeBaseModel):
9393
class BaseCreateInput(SafeBaseModel):
9494
name: str
95-
query: str | None = None
96-
repo_full_name: str | None = None
95+
query: str
96+
language: ProgrammingLanguage
9797

9898
input: BaseCreateInput
9999

@@ -102,8 +102,7 @@ class CreateResponse(SafeBaseModel):
102102
success: bool
103103
response: str
104104
code: str
105-
codemod_id: int
106-
context: str | None = None
105+
context: str
107106

108107

109108
###########################################################################

src/codegen/cli/auth/session.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from codegen.cli.errors import AuthError, NoTokenError
99
from codegen.cli.git.repo import get_git_repo
1010
from codegen.cli.utils.config import Config, get_config, write_config
11+
from codegen.sdk.enums import ProgrammingLanguage
1112

1213

1314
@dataclass
@@ -109,6 +110,13 @@ def repo_name(self) -> str:
109110
"""Get the current repository name"""
110111
return self.config.repo_full_name
111112

113+
@property
114+
def language(self) -> ProgrammingLanguage:
115+
"""Get the current language"""
116+
# TODO(jayhack): This is a temporary solution to get the language.
117+
# We should eventually get the language on init.
118+
return self.config.programming_language or ProgrammingLanguage.PYTHON
119+
112120
@property
113121
def codegen_dir(self) -> Path:
114122
"""Get the path to the codegen-sh directory"""

src/codegen/cli/codemod/convert.py

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,21 @@
22

33

44
def convert_to_cli(input: str, language: str, name: str) -> str:
5-
codebase_type = "PyCodebaseType" if language.lower() == "python" else "TSCodebaseType"
6-
return f"""import codegen.cli.sdk.decorator
7-
# from app.codemod.compilation.models.context import CodemodContext
8-
#from app.codemod.compilation.models.pr_options import PROptions
5+
return f"""import codegen
6+
from codegen import Codebase
97
10-
from codegen.sdk import {codebase_type}
118
12-
context: Any
9+
@codegen.function('{name}')
10+
def run(codebase: Codebase):
11+
{indent(input, " ")}
1312
1413
15-
@codegen.cli.sdk.decorator.function('{name}')
16-
def run(codebase: {codebase_type}, pr_options: Any):
17-
{indent(input, " ")}
14+
if __name__ == "__main__":
15+
print('Parsing codebase...')
16+
codebase = Codebase("./")
17+
18+
print('Running function...')
19+
codegen.run(run)
1820
"""
1921

2022

src/codegen/cli/commands/create/main.py

Lines changed: 15 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
from codegen.cli.rich.codeblocks import format_command, format_path
1212
from codegen.cli.rich.pretty_print import pretty_print_error
1313
from codegen.cli.rich.spinners import create_spinner
14-
from codegen.cli.utils.constants import ProgrammingLanguage
1514
from codegen.cli.utils.default_code import DEFAULT_CODEMOD
1615
from codegen.cli.workspace.decorators import requires_init
1716

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

3130

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

5253

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

8182
# Check if file exists
82-
if target_path.exists() and not overwrite:
83-
rel_path = make_relative(target_path)
83+
if codemod_path.exists() and not overwrite:
84+
rel_path = make_relative(codemod_path)
8485
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')}")
8586
return
8687

@@ -89,34 +90,30 @@ def create_command(session: CodegenSession, name: str, path: Path, description:
8990
try:
9091
if description:
9192
# Use API to generate implementation
92-
with create_spinner("Generating function (using LLM, this will take ~30s)") as status:
93+
with create_spinner("Generating function (using LLM, this will take ~10s)") as status:
9394
response = RestAPI(session.token).create(name=name, query=description)
94-
code = convert_to_cli(response.code, session.config.programming_language or ProgrammingLanguage.PYTHON, name)
95-
96-
# Write the system prompt if provided
97-
if response.context:
98-
prompt_path = get_prompts_dir() / f"{name.lower().replace(' ', '-')}-system-prompt.md"
99-
prompt_path.write_text(response.context)
95+
code = convert_to_cli(response.code, session.language, name)
10096
else:
10197
# Use default implementation
10298
code = DEFAULT_CODEMOD.format(name=name)
10399

104100
# Create the target directory if needed
105-
target_path.parent.mkdir(parents=True, exist_ok=True)
101+
codemod_path.parent.mkdir(parents=True, exist_ok=True)
106102

107103
# Write the function code
108-
target_path.write_text(code)
104+
codemod_path.write_text(code)
105+
prompt_path.write_text(response.context)
109106

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

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

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

0 commit comments

Comments
 (0)