Update ruff config and fix issues

ruff.toml

@@ -6,7 +6,20 @@ exclude = [
 ]
 unsafe-fixes = true
 [lint]
-select = ["F", "E", "W", "I", "UP", "D", "RUF", "DTZ", "TC", "EM", "TRY400"]
+select = [
+  "F",
+  "E",
+  "W",
+  "I",
+  "UP",
+  "D",
+  "YTT",
+  "RUF",
+  "DTZ",
+  "TC",
+  "EM",
+  "TRY400",
+]
 ignore = [
   "D100",
   "D101",
@@ -216,3 +229,5 @@ extend-generics = [
 ]
 [lint.isort]
 known-first-party = ["codegen"]
+[format]
+docstring-code-format = true

scripts/profiling/profile.py

@@ -5,7 +5,7 @@
 import typer
 
 
-def profile(repo: str, memory: bool = False):
+def profile(repo: str, memory: bool = False) -> None:
     type = "mem" if memory else "cpu"
     base = f".profiles/{type}/{repo}"
     os.makedirs(base, exist_ok=True)
@@ -37,7 +37,7 @@ def profile(repo: str, memory: bool = False):
         test.kill()
 
 
-def main():
+def main() -> None:
     typer.run(profile)
 
 

src/codegen/cli/api/client.py

@@ -55,7 +55,7 @@ class RestAPI:
 
     auth_token: str | None = None
 
-    def __init__(self, auth_token: str):
+    def __init__(self, auth_token: str) -> None:
         self.auth_token = auth_token
 
     def _get_headers(self) -> dict[str, str]:

src/codegen/cli/api/modal.py

@@ -2,7 +2,7 @@
 from codegen.cli.env.global_env import global_env
 
 
-def get_modal_workspace():
+def get_modal_workspace() -> str:
     match global_env.ENV:
         case Environment.PRODUCTION:
             return "codegen-sh"

src/codegen/cli/auth/session.py

@@ -46,7 +46,7 @@ class CodegenSession:
     _identity: Identity | None = None
     _profile: UserProfile | None = None
 
-    def __init__(self, token: str | None = None):
+    def __init__(self, token: str | None = None) -> None:
         self.token = token or get_current_token()
 
     @property

src/codegen/cli/auth/token_manager.py

@@ -9,12 +9,12 @@ class TokenManager:
     # Simple token manager to store and retrieve tokens.
     # This manager checks if the token is expired before retrieval.
     # TODO: add support for refreshing token and re authorization via supabase oauth
-    def __init__(self):
+    def __init__(self) -> None:
         self.config_dir = CONFIG_DIR
         self.token_file = AUTH_FILE
         self._ensure_config_dir()
 
-    def _ensure_config_dir(self):
+    def _ensure_config_dir(self) -> None:
         """Create config directory if it doesn't exist."""
         if not os.path.exists(self.config_dir):
             Path(self.config_dir).mkdir(parents=True, exist_ok=True)

src/codegen/cli/cli.py

@@ -21,7 +21,7 @@
 
 @click.group()
 @click.version_option(prog_name="codegen", message="%(version)s")
-def main():
+def main() -> None:
     """Codegen CLI - Transform your code with AI."""
 
 

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

@@ -70,7 +70,7 @@ def make_relative(path: Path) -> str:
 @click.argument("path", type=click.Path(path_type=Path), default=Path.cwd())
 @click.option("--description", "-d", default=None, help="Description of what this codemod does.")
 @click.option("--overwrite", is_flag=True, help="Overwrites function if it already exists.")
-def create_command(session: CodegenSession, name: str, path: Path, description: str | None = None, overwrite: bool = False):
+def create_command(session: CodegenSession, name: str, path: Path, description: str | None = None, overwrite: bool = False) -> None:
     """Create a new codegen function.
 
     NAME is the name/label for the function

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

@@ -47,7 +47,7 @@ def deploy_functions(session: CodegenSession, functions: list[DecoratedFunction]
 @click.argument("name", required=False)
 @click.option("-d", "--directory", type=click.Path(exists=True, path_type=Path), help="Directory to search for functions")
 @click.option("-m", "--message", help="Optional message to include with the deploy")
-def deploy_command(session: CodegenSession, name: str | None = None, directory: Path | None = None, message: str | None = None):
+def deploy_command(session: CodegenSession, name: str | None = None, directory: Path | None = None, message: str | None = None) -> None:
     """Deploy codegen functions.
 
     If NAME is provided, deploys a specific function by that name.

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

@@ -13,7 +13,7 @@
 @click.option("--query", "-q", help="The question to ask the expert.")
 @requires_auth
 @requires_init
-def expert_command(session: CodegenSession, query: str):
+def expert_command(session: CodegenSession, query: str) -> None:
     """Asks a codegen expert a question."""
     status = Status("Asking expert...", spinner="dots", spinner_style="purple")
     status.start()

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

@@ -17,7 +17,7 @@
 @click.option("--repo-name", type=str, help="The name of the repository")
 @click.option("--organization-name", type=str, help="The name of the organization")
 @click.option("--fetch-docs", is_flag=True, help="Fetch docs and examples (requires auth)")
-def init_command(repo_name: str | None = None, organization_name: str | None = None, fetch_docs: bool = False):
+def init_command(repo_name: str | None = None, organization_name: str | None = None, fetch_docs: bool = False) -> None:
     """Initialize or update the Codegen folder."""
     # Print a message if not in a git repo
     try:

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

@@ -9,7 +9,7 @@
 
 
 @click.command(name="list")
-def list_command():
+def list_command() -> None:
     """List available codegen functions."""
     functions = CodemodManager.get_decorated()
     if functions:

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

@@ -8,7 +8,7 @@
 
 @click.command(name="login")
 @click.option("--token", required=False, help="API token for authentication")
-def login_command(token: str):
+def login_command(token: str) -> None:
     """Store authentication token."""
     # Check if already authenticated
     token_manager = TokenManager()

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

@@ -5,7 +5,7 @@
 
 
 @click.command(name="logout")
-def logout_command():
+def logout_command() -> None:
     """Clear stored authentication token."""
     token_manager = TokenManager()
     token_manager.clear_token()

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

@@ -23,7 +23,7 @@ def create_jupyter_dir() -> Path:
 @click.option("--background", is_flag=True, help="Run Jupyter Lab in the background")
 @click.option("--demo", is_flag=True, help="Create a demo notebook with FastAPI example code")
 @requires_init
-def notebook_command(session: CodegenSession, background: bool, demo: bool):
+def notebook_command(session: CodegenSession, background: bool, demo: bool) -> None:
     """Launch Jupyter Lab with a pre-configured notebook for exploring your codebase."""
     with create_spinner("Setting up Jupyter environment...") as status:
         venv = VenvManager()

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

@@ -11,7 +11,7 @@
 @click.command(name="profile")
 @requires_auth
 @requires_init
-def profile_command(session: CodegenSession):
+def profile_command(session: CodegenSession) -> None:
     """Display information about the currently authenticated user."""
     rich.print(
         Panel(

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

@@ -22,7 +22,7 @@ def run_command(
     web: bool = False,
     diff_preview: int | None = None,
     arguments: str | None = None,
-):
+) -> None:
     """Run a codegen function by its label."""
     # Ensure venv is initialized
     venv = VenvManager()

src/codegen/cli/commands/run/render.py

@@ -6,7 +6,7 @@
 from codegen.cli.api.schemas import RunCodemodOutput
 
 
-def pretty_print_output(output: RunCodemodOutput):
+def pretty_print_output(output: RunCodemodOutput) -> None:
     """Pretty print the codemod run output with panels."""
     if output.web_link:
         rich.print("\n• [blue underline]" + output.web_link + "[/blue underline]\n")
@@ -21,7 +21,7 @@ def pretty_print_output(output: RunCodemodOutput):
         pretty_print_diff(output.observation)
 
 
-def pretty_print_logs(logs: str):
+def pretty_print_logs(logs: str) -> None:
     """Pretty print logs in a panel."""
     rich.print(
         Panel(
@@ -35,7 +35,7 @@ def pretty_print_logs(logs: str):
     rich.print()  # spacing
 
 
-def pretty_print_error(error: str):
+def pretty_print_error(error: str) -> None:
     """Pretty print error in a panel."""
     rich.print(
         Panel(
@@ -49,7 +49,7 @@ def pretty_print_error(error: str):
     rich.print()  # spacing
 
 
-def pretty_print_diff(diff: str):
+def pretty_print_diff(diff: str) -> None:
     """Pretty print diff in a panel."""
     rich.print(
         Panel(

src/codegen/cli/commands/run/run_cloud.py

@@ -13,7 +13,7 @@
 from codegen.cli.utils.url import generate_webapp_url
 
 
-def run_cloud(session: CodegenSession, function, apply_local: bool = False, diff_preview: int | None = None):
+def run_cloud(session: CodegenSession, function, apply_local: bool = False, diff_preview: int | None = None) -> None:
     """Run a function on the cloud service.
 
     Args:

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

@@ -36,7 +36,7 @@ def run_on_pr(session: CodegenSession, codemod_name: str, pr_number: int) -> Non
 @requires_auth
 @click.argument("codemod_name", type=str)
 @click.argument("pr_number", type=int)
-def run_on_pr_command(session: CodegenSession, codemod_name: str, pr_number: int):
+def run_on_pr_command(session: CodegenSession, codemod_name: str, pr_number: int) -> None:
     """Test a webhook against a specific PR.
 
     CODEMOD_NAME is the name of the codemod to test

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

@@ -9,7 +9,7 @@
 
 @click.command(name="style-debug")
 @click.option("--text", default="Loading...", help="Text to show in the spinner")
-def style_debug_command(text: str):
+def style_debug_command(text: str) -> None:
     """Debug command to visualize CLI styling (spinners, etc)."""
     try:
         with create_spinner(text) as status:

src/codegen/cli/errors.py

@@ -36,7 +36,7 @@ class ServerError(CodegenError):
     pass
 
 
-def format_error_message(error):
+def format_error_message(error) -> str:
     """Format error message based on error type."""
     if isinstance(error, AuthError):
         return "[red]Authentication Error:[/red] Please run 'codegen login' first."

src/codegen/cli/git/patch.py

@@ -4,7 +4,7 @@
 from rich.status import Status
 
 
-def apply_patch(git_repo: Repository, patch: str):
+def apply_patch(git_repo: Repository, patch: str) -> None:
     """Apply a git patch to the repository.
 
     Args:

src/codegen/cli/rich/pretty_print.py

@@ -6,7 +6,7 @@
 from codegen.cli.api.schemas import RunCodemodOutput
 
 
-def pretty_print_output(output: RunCodemodOutput):
+def pretty_print_output(output: RunCodemodOutput) -> None:
     """Pretty print the codemod run output with panels."""
     if output.web_link:
         rich.print("\n• [blue underline]" + output.web_link + "[/blue underline]\n")
@@ -21,7 +21,7 @@ def pretty_print_output(output: RunCodemodOutput):
         pretty_print_diff(output.observation)
 
 
-def pretty_print_logs(logs: str):
+def pretty_print_logs(logs: str) -> None:
     """Pretty print logs in a panel."""
     rich.print(
         Panel(
@@ -35,7 +35,7 @@ def pretty_print_logs(logs: str):
     rich.print()  # spacing
 
 
-def pretty_print_error(error: str):
+def pretty_print_error(error: str) -> None:
     """Pretty print error in a panel."""
     rich.print(
         Panel(
@@ -49,7 +49,7 @@ def pretty_print_error(error: str):
     rich.print()  # spacing
 
 
-def pretty_print_diff(diff: str):
+def pretty_print_diff(diff: str) -> None:
     """Pretty print diff in a panel."""
     rich.print(
         Panel(

src/codegen/cli/sdk/decorator.py

@@ -18,7 +18,7 @@ def __init__(
         webhook_config: dict | None = None,
         lint_mode: bool = False,
         lint_user_whitelist: Sequence[str] | None = None,
-    ):
+    ) -> None:
         self.name = name
         self.func: Callable | None = None
         self.params_type = None

src/codegen/cli/sdk/pull_request.py

@@ -14,7 +14,7 @@ class CodegenPullRequest:
     codegen_pr_id: int
     patch_set: PatchSet | None = None
 
-    def __init__(self, url: str, number: int, title: str, github_pr_number: int, codegen_pr_id: int, patch_set: PatchSet | None = None):
+    def __init__(self, url: str, number: int, title: str, github_pr_number: int, codegen_pr_id: int, patch_set: PatchSet | None = None) -> None:
         self.url = url
         self.number = number
         self.title = title
@@ -47,5 +47,5 @@ def lookup(cls, number: int) -> "CodegenPullRequest":
             patch_set=None,  # Can be loaded on demand if needed
         )
 
-    def __str__(self):
+    def __str__(self) -> str:
         return f"CodegenPullRequest(url={self.url}, number={self.number}, title={self.title}, github_pr_number={self.github_pr_number}, codegen_pr_id={self.codegen_pr_id})"

src/codegen/cli/utils/count_functions.py

@@ -20,7 +20,7 @@ class CountFunctionsArgs(BaseModel):
 
 
 @codegen.cli.sdk.decorator.function("count-functions")
-def run(codebase, pr_options, arguments: CountFunctionsArgs):
+def run(codebase, pr_options, arguments: CountFunctionsArgs) -> None:
     # Count Functions in Codebase
 
     # Initialize a total function counter

src/codegen/cli/utils/function_finder.py

@@ -80,7 +80,7 @@ def validate(self) -> None:
 
 
 class CodegenFunctionVisitor(ast.NodeVisitor):
-    def __init__(self):
+    def __init__(self) -> None:
         self.functions: list[DecoratedFunction] = []
 
     def get_function_body(self, node: ast.FunctionDef) -> str:
@@ -174,7 +174,7 @@ def get_function_parameters(self, node: ast.FunctionDef) -> list[tuple[str, str
 
         return parameters
 
-    def visit_FunctionDef(self, node):
+    def visit_FunctionDef(self, node) -> None:
         for decorator in node.decorator_list:
             if (
                 isinstance(decorator, ast.Call)
@@ -219,7 +219,7 @@ def _get_decorator_attrs(self, node):
             node = node.value
         return attrs
 
-    def visit_Module(self, node):
+    def visit_Module(self, node) -> None:
         # Store the full source code for later use
         self.source = self.file_content
         self.generic_visit(node)

src/codegen/cli/workspace/docs_workspace.py

@@ -4,7 +4,7 @@
 from rich.status import Status
 
 
-def populate_api_docs(dest: Path, api_docs: dict[str, str], status: Status):
+def populate_api_docs(dest: Path, api_docs: dict[str, str], status: Status) -> None:
     """Writes all API docs to the docs folder"""
     status.update("Populating API documentation...")
     # Remove existing docs

src/codegen/cli/workspace/examples_workspace.py

@@ -8,7 +8,7 @@
 from codegen.cli.codemod.convert import convert_to_cli
 
 
-def populate_examples(session: CodegenSession, dest: Path, examples: list[SerializedExample], status: Status):
+def populate_examples(session: CodegenSession, dest: Path, examples: list[SerializedExample], status: Status) -> None:
     """Populate the examples folder with examples for the current repository."""
     status.update("Populating example codemods...")
     # Remove existing examples

src/codegen/cli/workspace/initialize_workspace.py

@@ -117,14 +117,14 @@ def initialize_codegen(
     return CODEGEN_FOLDER, DOCS_FOLDER, EXAMPLES_FOLDER
 
 
-def add_to_gitignore_if_not_present(gitignore: Path, line: str):
+def add_to_gitignore_if_not_present(gitignore: Path, line: str) -> None:
     if not gitignore.exists():
         gitignore.write_text(line)
     elif line not in gitignore.read_text():
         gitignore.write_text(gitignore.read_text() + "\n" + line)
 
 
-def modify_gitignore(codegen_folder: Path):
+def modify_gitignore(codegen_folder: Path) -> None:
     """Update .gitignore to track only specific Codegen files."""
     gitignore_path = codegen_folder / ".gitignore"