universal-mcp 0.1.15rc7__py3-none-any.whl → 0.1.16__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -115,9 +115,7 @@ class ToolManager:
115
115
  f"Tool name '{tool.name}' conflicts with an existing tool. Skipping addition of new function."
116
116
  )
117
117
  else:
118
- logger.debug(
119
- f"Tool '{tool.name}' with the same function already exists."
120
- )
118
+ logger.debug(f"Tool '{tool.name}' with the same function already exists.")
121
119
  return existing
122
120
 
123
121
  logger.debug(f"Adding tool: {tool.name}")
@@ -170,9 +168,7 @@ class ToolManager:
170
168
  return
171
169
 
172
170
  if not isinstance(functions, list):
173
- logger.error(
174
- f"App '{app.name}' list_tools() did not return a list. Skipping registration."
175
- )
171
+ logger.error(f"App '{app.name}' list_tools() did not return a list. Skipping registration.")
176
172
  return
177
173
 
178
174
  tools = []
@@ -183,18 +179,12 @@ class ToolManager:
183
179
 
184
180
  try:
185
181
  tool_instance = Tool.from_function(function)
186
- tool_instance.name = (
187
- f"{app.name}{TOOL_NAME_SEPARATOR}{tool_instance.name}"
188
- )
189
- tool_instance.tags.append(
190
- app.name
191
- ) if app.name not in tool_instance.tags else None
182
+ tool_instance.name = f"{app.name}{TOOL_NAME_SEPARATOR}{tool_instance.name}"
183
+ tool_instance.tags.append(app.name) if app.name not in tool_instance.tags else None
192
184
  tools.append(tool_instance)
193
185
  except Exception as e:
194
186
  tool_name = getattr(function, "__name__", "unknown")
195
- logger.error(
196
- f"Failed to create Tool from '{tool_name}' in {app.name}: {e}"
197
- )
187
+ logger.error(f"Failed to create Tool from '{tool_name}' in {app.name}: {e}")
198
188
 
199
189
  tools = _filter_by_name(tools, tool_names)
200
190
  tools = _filter_by_tags(tools, tags)
@@ -20,20 +20,15 @@ class Tool(BaseModel):
20
20
  args_description: dict[str, str] = Field(
21
21
  default_factory=dict, description="Descriptions of arguments from the docstring"
22
22
  )
23
- returns_description: str = Field(
24
- default="", description="Description of the return value from the docstring"
25
- )
23
+ returns_description: str = Field(default="", description="Description of the return value from the docstring")
26
24
  raises_description: dict[str, str] = Field(
27
25
  default_factory=dict,
28
26
  description="Descriptions of exceptions raised from the docstring",
29
27
  )
30
- tags: list[str] = Field(
31
- default_factory=list, description="Tags for categorizing the tool"
32
- )
28
+ tags: list[str] = Field(default_factory=list, description="Tags for categorizing the tool")
33
29
  parameters: dict[str, Any] = Field(description="JSON schema for tool parameters")
34
30
  fn_metadata: FuncMetadata = Field(
35
- description="Metadata about the function including a pydantic model for tool"
36
- " arguments"
31
+ description="Metadata about the function including a pydantic model for tool arguments"
37
32
  )
38
33
  is_async: bool = Field(description="Whether the tool is async")
39
34
 
@@ -55,9 +50,7 @@ class Tool(BaseModel):
55
50
 
56
51
  is_async = inspect.iscoroutinefunction(fn)
57
52
 
58
- func_arg_metadata = FuncMetadata.func_metadata(
59
- fn, arg_description=parsed_doc["args"]
60
- )
53
+ func_arg_metadata = FuncMetadata.func_metadata(fn, arg_description=parsed_doc["args"])
61
54
  parameters = func_arg_metadata.arg_model.model_json_schema()
62
55
 
63
56
  return cls(
@@ -26,9 +26,7 @@ class AgentrClient(metaclass=Singleton):
26
26
  "API key for AgentR is missing. Please visit https://agentr.dev to create an API key, then set it as AGENTR_API_KEY environment variable."
27
27
  )
28
28
  raise ValueError("AgentR API key required - get one at https://agentr.dev")
29
- self.base_url = (
30
- base_url or os.getenv("AGENTR_BASE_URL", "https://api.agentr.dev")
31
- ).rstrip("/")
29
+ self.base_url = (base_url or os.getenv("AGENTR_BASE_URL", "https://api.agentr.dev")).rstrip("/")
32
30
 
33
31
  def get_credentials(self, integration_name: str) -> dict:
34
32
  """Get credentials for an integration from the AgentR API.
@@ -48,9 +46,7 @@ class AgentrClient(metaclass=Singleton):
48
46
  headers={"accept": "application/json", "X-API-KEY": self.api_key},
49
47
  )
50
48
  if response.status_code == 404:
51
- logger.warning(
52
- f"No credentials found for {integration_name}. Requesting authorization..."
53
- )
49
+ logger.warning(f"No credentials found for {integration_name}. Requesting authorization...")
54
50
  action = self.get_authorization_url(integration_name)
55
51
  raise NotAuthorizedError(action)
56
52
  response.raise_for_status()
@@ -3,7 +3,7 @@ def get_default_repository_path(slug: str) -> str:
3
3
  Convert a repository slug to a repository URL.
4
4
  """
5
5
  slug = slug.strip().lower()
6
- return f"git+https://github.com/universal-mcp/{slug}"
6
+ return f"universal-mcp-{slug}"
7
7
 
8
8
 
9
9
  def get_default_package_name(slug: str) -> str:
@@ -81,17 +81,13 @@ def parse_docstring(docstring: str | None) -> dict[str, Any]:
81
81
  elif stripped_lower.startswith(("raises ", "errors ", "exceptions ")):
82
82
  section_type = "raises"
83
83
  # Capture content after header word and potential colon/space
84
- parts = re.split(
85
- r"[:\s]+", line.strip(), maxsplit=1
86
- ) # B034: Use keyword maxsplit
84
+ parts = re.split(r"[:\s]+", line.strip(), maxsplit=1) # B034: Use keyword maxsplit
87
85
  if len(parts) > 1:
88
86
  header_content = parts[1].strip()
89
87
  elif stripped_lower.startswith(("tags",)):
90
88
  section_type = "tags"
91
89
  # Capture content after header word and potential colon/space
92
- parts = re.split(
93
- r"[:\s]+", line.strip(), maxsplit=1
94
- ) # B034: Use keyword maxsplit
90
+ parts = re.split(r"[:\s]+", line.strip(), maxsplit=1) # B034: Use keyword maxsplit
95
91
  if len(parts) > 1:
96
92
  header_content = parts[1].strip()
97
93
 
@@ -117,9 +113,7 @@ def parse_docstring(docstring: str | None) -> dict[str, Any]:
117
113
  stripped_line = line.strip()
118
114
  original_indentation = len(line) - len(line.lstrip(" "))
119
115
 
120
- is_new_section_header, new_section_type_this_line, header_content_this_line = (
121
- check_for_section_header(line)
122
- )
116
+ is_new_section_header, new_section_type_this_line, header_content_this_line = check_for_section_header(line)
123
117
 
124
118
  should_finalize_previous = False
125
119
 
@@ -156,10 +150,7 @@ def parse_docstring(docstring: str | None) -> dict[str, Any]:
156
150
  or (
157
151
  current_section in ["args", "raises"]
158
152
  and current_key is not None
159
- and (
160
- key_pattern.match(line)
161
- or (original_indentation == 0 and stripped_line)
162
- )
153
+ and (key_pattern.match(line) or (original_indentation == 0 and stripped_line))
163
154
  )
164
155
  or (
165
156
  current_section in ["returns", "tags", "other"]
@@ -7,19 +7,30 @@ from loguru import logger
7
7
  from rich import print
8
8
 
9
9
 
10
- def get_uvx_path() -> str:
10
+ def get_uvx_path() -> Path:
11
11
  """Get the full path to the uv executable."""
12
12
  uvx_path = shutil.which("uvx")
13
- if not uvx_path:
14
- logger.error(
15
- "uvx executable not found in PATH, falling back to 'uvx'. "
16
- "Please ensure uvx is installed and in your PATH"
17
- )
18
- return "uvx" # Fall back to just "uvx" if not found
19
- return uvx_path
20
-
21
-
22
- def create_file_if_not_exists(path: Path) -> None:
13
+ if uvx_path:
14
+ return Path(uvx_path)
15
+
16
+ # Check if
17
+ common_uv_paths = [
18
+ Path.home() / ".local/bin/uvx", # Linux/macOS
19
+ Path.home() / ".cargo/bin/uvx", # Linux/macOS
20
+ Path.home() / "AppData/Local/Programs/uvx/uvx.exe", # Windows
21
+ Path.home() / "AppData/Roaming/uvx/uvx.exe", # Windows
22
+ ]
23
+ for path in common_uv_paths:
24
+ logger.info(f"Checking {path}")
25
+ if path.exists():
26
+ return path
27
+ logger.error(
28
+ "uvx executable not found in PATH, falling back to 'uvx'. Please ensure uvx is installed and in your PATH"
29
+ )
30
+ return None # Fall back to just "uvx" if not found
31
+
32
+
33
+ def _create_file_if_not_exists(path: Path) -> None:
23
34
  """Create a file if it doesn't exist"""
24
35
  if not path.exists():
25
36
  print(f"[yellow]Creating config file at {path}[/yellow]")
@@ -27,6 +38,34 @@ def create_file_if_not_exists(path: Path) -> None:
27
38
  json.dump({}, f)
28
39
 
29
40
 
41
+ def _generate_mcp_config(api_key: str) -> None:
42
+ uvx_path = get_uvx_path()
43
+ if not uvx_path:
44
+ raise ValueError("uvx executable not found in PATH")
45
+ return {
46
+ "command": str(uvx_path),
47
+ "args": ["universal_mcp@latest", "run"],
48
+ "env": {
49
+ "AGENTR_API_KEY": api_key,
50
+ "PATH": str(uvx_path.parent),
51
+ },
52
+ }
53
+
54
+
55
+ def _install_config(config_path: Path, mcp_config: dict):
56
+ _create_file_if_not_exists(config_path)
57
+ try:
58
+ config = json.loads(config_path.read_text())
59
+ except json.JSONDecodeError:
60
+ print("[yellow]Config file was empty or invalid, creating new configuration[/yellow]")
61
+ config = {}
62
+ if "mcpServers" not in config:
63
+ config["mcpServers"] = {}
64
+ config["mcpServers"]["universal_mcp"] = mcp_config
65
+ with open(config_path, "w") as f:
66
+ json.dump(config, f, indent=4)
67
+
68
+
30
69
  def get_supported_apps() -> list[str]:
31
70
  """Get list of supported apps"""
32
71
  return ["claude", "cursor", "cline", "continue", "goose", "windsurf", "zed"]
@@ -37,35 +76,17 @@ def install_claude(api_key: str) -> None:
37
76
  print("[bold blue]Installing Claude configuration...[/bold blue]")
38
77
  # Determine platform-specific config path
39
78
  if sys.platform == "darwin": # macOS
40
- config_path = (
41
- Path.home()
42
- / "Library/Application Support/Claude/claude_desktop_config.json"
43
- )
79
+ config_path = Path.home() / "Library/Application Support/Claude/claude_desktop_config.json"
44
80
  elif sys.platform == "win32": # Windows
45
81
  config_path = Path.home() / "AppData/Roaming/Claude/claude_desktop_config.json"
46
82
  else:
47
83
  raise ValueError(
48
84
  "Unsupported platform. Only macOS and Windows are currently supported.",
49
85
  )
50
-
51
- # Create config directory if it doesn't exist
52
- create_file_if_not_exists(config_path)
53
- try:
54
- config = json.loads(config_path.read_text())
55
- except json.JSONDecodeError:
56
- print(
57
- "[yellow]Config file was empty or invalid, creating new configuration[/yellow]"
58
- )
59
- config = {}
60
- if "mcpServers" not in config:
61
- config["mcpServers"] = {}
62
- config["mcpServers"]["universal_mcp"] = {
63
- "command": get_uvx_path(),
64
- "args": ["universal_mcp@latest", "run"],
65
- "env": {"AGENTR_API_KEY": api_key},
66
- }
67
- with open(config_path, "w") as f:
68
- json.dump(config, f, indent=4)
86
+ mcp_config = _generate_mcp_config(api_key=api_key)
87
+ logger.info(f"Installing Claude configuration at {config_path}")
88
+ logger.info(f"Config: {mcp_config}")
89
+ _install_config(config_path=config_path, mcp_config=mcp_config)
69
90
  print("[green]✓[/green] Claude configuration installed successfully")
70
91
 
71
92
 
@@ -74,28 +95,8 @@ def install_cursor(api_key: str) -> None:
74
95
  print("[bold blue]Installing Cursor configuration...[/bold blue]")
75
96
  # Set up Cursor config path
76
97
  config_path = Path.home() / ".cursor/mcp.json"
77
-
78
- # Create config directory if it doesn't exist
79
- create_file_if_not_exists(config_path)
80
-
81
- try:
82
- config = json.loads(config_path.read_text())
83
- except json.JSONDecodeError:
84
- print(
85
- "[yellow]Config file was empty or invalid, creating new configuration[/yellow]"
86
- )
87
- config = {}
88
-
89
- if "mcpServers" not in config:
90
- config["mcpServers"] = {}
91
- config["mcpServers"]["universal_mcp"] = {
92
- "command": get_uvx_path(),
93
- "args": ["universal_mcp@latest", "run"],
94
- "env": {"AGENTR_API_KEY": api_key},
95
- }
96
-
97
- with open(config_path, "w") as f:
98
- json.dump(config, f, indent=4)
98
+ mcp_config = _generate_mcp_config(api_key=api_key)
99
+ _install_config(config_path=config_path, mcp_config=mcp_config)
99
100
  print("[green]✓[/green] Cursor configuration installed successfully")
100
101
 
101
102
 
@@ -104,28 +105,8 @@ def install_cline(api_key: str) -> None:
104
105
  print("[bold blue]Installing Cline configuration...[/bold blue]")
105
106
  # Set up Cline config path
106
107
  config_path = Path.home() / ".config/cline/mcp.json"
107
-
108
- # Create config directory if it doesn't exist
109
- create_file_if_not_exists(config_path)
110
-
111
- try:
112
- config = json.loads(config_path.read_text())
113
- except json.JSONDecodeError:
114
- print(
115
- "[yellow]Config file was empty or invalid, creating new configuration[/yellow]"
116
- )
117
- config = {}
118
-
119
- if "mcpServers" not in config:
120
- config["mcpServers"] = {}
121
- config["mcpServers"]["universal_mcp"] = {
122
- "command": get_uvx_path(),
123
- "args": ["universal_mcp@latest", "run"],
124
- "env": {"AGENTR_API_KEY": api_key},
125
- }
126
-
127
- with open(config_path, "w") as f:
128
- json.dump(config, f, indent=4)
108
+ mcp_config = _generate_mcp_config(api_key=api_key)
109
+ _install_config(config_path=config_path, mcp_config=mcp_config)
129
110
  print("[green]✓[/green] Cline configuration installed successfully")
130
111
 
131
112
 
@@ -140,28 +121,8 @@ def install_continue(api_key: str) -> None:
140
121
  config_path = Path.home() / "AppData/Roaming/Continue/mcp.json"
141
122
  else: # Linux and others
142
123
  config_path = Path.home() / ".config/continue/mcp.json"
143
-
144
- # Create config directory if it doesn't exist
145
- create_file_if_not_exists(config_path)
146
-
147
- try:
148
- config = json.loads(config_path.read_text())
149
- except json.JSONDecodeError:
150
- print(
151
- "[yellow]Config file was empty or invalid, creating new configuration[/yellow]"
152
- )
153
- config = {}
154
-
155
- if "mcpServers" not in config:
156
- config["mcpServers"] = {}
157
- config["mcpServers"]["universal_mcp"] = {
158
- "command": get_uvx_path(),
159
- "args": ["universal_mcp@latest", "run"],
160
- "env": {"AGENTR_API_KEY": api_key},
161
- }
162
-
163
- with open(config_path, "w") as f:
164
- json.dump(config, f, indent=4)
124
+ mcp_config = _generate_mcp_config(api_key=api_key)
125
+ _install_config(config_path=config_path, mcp_config=mcp_config)
165
126
  print("[green]✓[/green] Continue configuration installed successfully")
166
127
 
167
128
 
@@ -177,27 +138,8 @@ def install_goose(api_key: str) -> None:
177
138
  else: # Linux and others
178
139
  config_path = Path.home() / ".config/goose/mcp-config.json"
179
140
 
180
- # Create config directory if it doesn't exist
181
- create_file_if_not_exists(config_path)
182
-
183
- try:
184
- config = json.loads(config_path.read_text())
185
- except json.JSONDecodeError:
186
- print(
187
- "[yellow]Config file was empty or invalid, creating new configuration[/yellow]"
188
- )
189
- config = {}
190
-
191
- if "mcpServers" not in config:
192
- config["mcpServers"] = {}
193
- config["mcpServers"]["universal_mcp"] = {
194
- "command": get_uvx_path(),
195
- "args": ["universal_mcp@latest", "run"],
196
- "env": {"AGENTR_API_KEY": api_key},
197
- }
198
-
199
- with open(config_path, "w") as f:
200
- json.dump(config, f, indent=4)
141
+ mcp_config = _generate_mcp_config(api_key=api_key)
142
+ _install_config(config_path=config_path, mcp_config=mcp_config)
201
143
  print("[green]✓[/green] Goose configuration installed successfully")
202
144
 
203
145
 
@@ -212,28 +154,8 @@ def install_windsurf(api_key: str) -> None:
212
154
  config_path = Path.home() / "AppData/Roaming/Windsurf/mcp.json"
213
155
  else: # Linux and others
214
156
  config_path = Path.home() / ".config/windsurf/mcp.json"
215
-
216
- # Create config directory if it doesn't exist
217
- create_file_if_not_exists(config_path)
218
-
219
- try:
220
- config = json.loads(config_path.read_text())
221
- except json.JSONDecodeError:
222
- print(
223
- "[yellow]Config file was empty or invalid, creating new configuration[/yellow]"
224
- )
225
- config = {}
226
-
227
- if "mcpServers" not in config:
228
- config["mcpServers"] = {}
229
- config["mcpServers"]["universal_mcp"] = {
230
- "command": get_uvx_path(),
231
- "args": ["universal_mcp@latest", "run"],
232
- "env": {"AGENTR_API_KEY": api_key},
233
- }
234
-
235
- with open(config_path, "w") as f:
236
- json.dump(config, f, indent=4)
157
+ mcp_config = _generate_mcp_config(api_key=api_key)
158
+ _install_config(config_path=config_path, mcp_config=mcp_config)
237
159
  print("[green]✓[/green] Windsurf configuration installed successfully")
238
160
 
239
161
 
@@ -244,47 +166,8 @@ def install_zed(api_key: str) -> None:
244
166
  # Set up Zed config path
245
167
  config_dir = Path.home() / ".config/zed"
246
168
  config_path = config_dir / "mcp_servers.json"
247
-
248
- # Create config directory if it doesn't exist
249
- create_file_if_not_exists(config_path)
250
-
251
- try:
252
- config = json.loads(config_path.read_text())
253
- except json.JSONDecodeError:
254
- print(
255
- "[yellow]Config file was empty or invalid, creating new configuration[/yellow]"
256
- )
257
- config = {}
258
-
259
- if not isinstance(config, list):
260
- config = []
261
-
262
- # Check if universal_mcp is already in the config
263
- existing_config = False
264
- for server in config:
265
- if server.get("name") == "universal_mcp":
266
- existing_config = True
267
- server.update(
268
- {
269
- "command": get_uvx_path(),
270
- "args": ["universal_mcp@latest", "run"],
271
- "env": {"AGENTR_API_KEY": api_key},
272
- }
273
- )
274
- break
275
-
276
- if not existing_config:
277
- config.append(
278
- {
279
- "name": "universal_mcp",
280
- "command": get_uvx_path(),
281
- "args": ["universal_mcp@latest", "run"],
282
- "env": {"AGENTR_API_KEY": api_key},
283
- }
284
- )
285
-
286
- with open(config_path, "w") as f:
287
- json.dump(config, f, indent=4)
169
+ mcp_config = _generate_mcp_config(api_key=api_key)
170
+ _install_config(config_path=config_path, mcp_config=mcp_config)
288
171
  print("[green]✓[/green] Zed configuration installed successfully")
289
172
 
290
173
 
@@ -109,9 +109,7 @@ def generate_api_from_schema(
109
109
  f.write(code)
110
110
 
111
111
  if not test_correct_output(gen_file):
112
- logger.error(
113
- "Generated code validation failed for '%s'. Aborting generation.", gen_file
114
- )
112
+ logger.error("Generated code validation failed for '%s'. Aborting generation.", gen_file)
115
113
  logger.info("Next steps:")
116
114
  logger.info(" 1) Review your OpenAPI schema for potential mismatches.")
117
115
  logger.info(
@@ -19,12 +19,8 @@ from pydantic import BaseModel, Field
19
19
  class DocstringOutput(BaseModel):
20
20
  """Structure for the generated docstring output."""
21
21
 
22
- summary: str = Field(
23
- description="A clear, concise summary of what the function does"
24
- )
25
- args: dict[str, str] = Field(
26
- description="Dictionary mapping parameter names to their descriptions"
27
- )
22
+ summary: str = Field(description="A clear, concise summary of what the function does")
23
+ args: dict[str, str] = Field(description="Dictionary mapping parameter names to their descriptions")
28
24
  returns: str = Field(description="Description of what the function returns")
29
25
  raises: dict[str, str] = Field(
30
26
  default_factory=dict,
@@ -44,9 +40,7 @@ class FunctionExtractor(ast.NodeVisitor):
44
40
 
45
41
  def __init__(self, source_code: str):
46
42
  self.source_lines = source_code.splitlines(keepends=True)
47
- self.functions: list[
48
- tuple[str, str]
49
- ] = [] # Store tuples of (function_name, function_source)
43
+ self.functions: list[tuple[str, str]] = [] # Store tuples of (function_name, function_source)
50
44
 
51
45
  def _get_source_segment(self, node: ast.AST) -> str | None:
52
46
  """Safely extracts the source segment for a node using ast.get_source_segment."""
@@ -110,9 +104,7 @@ def extract_functions_from_script(file_path: str) -> list[tuple[str, str]]:
110
104
  try:
111
105
  tree = ast.parse(source_code, filename=file_path)
112
106
  except SyntaxError as e:
113
- print(
114
- f"Error: Invalid Python syntax in {file_path} at line {e.lineno}, offset {e.offset}: {e.msg}"
115
- )
107
+ print(f"Error: Invalid Python syntax in {file_path} at line {e.lineno}, offset {e.offset}: {e.msg}")
116
108
  raise
117
109
  except Exception as e:
118
110
  print(f"Error parsing {file_path} into AST: {e}")
@@ -175,9 +167,7 @@ def extract_json_from_text(text):
175
167
  raise ValueError("Could not extract valid JSON from the response") from e
176
168
 
177
169
 
178
- def generate_docstring(
179
- function_code: str, model: str = "perplexity/sonar"
180
- ) -> DocstringOutput:
170
+ def generate_docstring(function_code: str, model: str = "perplexity/sonar") -> DocstringOutput:
181
171
  """
182
172
  Generate a docstring for a Python function using litellm with structured output.
183
173
 
@@ -239,9 +229,7 @@ def generate_docstring(
239
229
  parsed_data = extract_json_from_text(response_text)
240
230
  except ValueError as e:
241
231
  print(f"JSON extraction failed: {e}")
242
- print(
243
- f"Raw response: {response_text[:100]}..."
244
- ) # Log first 100 chars for debugging
232
+ print(f"Raw response: {response_text[:100]}...") # Log first 100 chars for debugging
245
233
  # Return a default structure if extraction fails
246
234
  return DocstringOutput(
247
235
  summary="Failed to extract docstring information",
@@ -291,20 +279,14 @@ def format_docstring(docstring: DocstringOutput) -> str:
291
279
  if summary:
292
280
  parts.append(summary)
293
281
 
294
- filtered_args = {
295
- name: desc
296
- for name, desc in docstring.args.items()
297
- if name not in ("self", "cls")
298
- }
282
+ filtered_args = {name: desc for name, desc in docstring.args.items() if name not in ("self", "cls")}
299
283
  args_lines = []
300
284
  if filtered_args:
301
285
  args_lines.append("Args:")
302
286
  for arg_name, arg_desc in filtered_args.items():
303
287
  arg_desc_cleaned = arg_desc.strip()
304
288
  args_lines.append(f" {arg_name}: {arg_desc_cleaned}")
305
- elif docstring.args.get(
306
- "None"
307
- ): # Include the 'None' placeholder if it was generated
289
+ elif docstring.args.get("None"): # Include the 'None' placeholder if it was generated
308
290
  args_lines.append("Args:")
309
291
  none_desc_cleaned = docstring.args["None"].strip()
310
292
  args_lines.append(f" None: {none_desc_cleaned}")
@@ -321,12 +303,8 @@ def format_docstring(docstring: DocstringOutput) -> str:
321
303
  raises_lines.append("Raises:")
322
304
  for exception_type, exception_desc in docstring.raises.items():
323
305
  exception_desc_cleaned = exception_desc.strip()
324
- if (
325
- exception_type.strip() and exception_desc_cleaned
326
- ): # Ensure type and desc are not empty
327
- raises_lines.append(
328
- f" {exception_type.strip()}: {exception_desc_cleaned}"
329
- )
306
+ if exception_type.strip() and exception_desc_cleaned: # Ensure type and desc are not empty
307
+ raises_lines.append(f" {exception_type.strip()}: {exception_desc_cleaned}")
330
308
  if raises_lines:
331
309
  parts.append("\n".join(raises_lines))
332
310
 
@@ -363,9 +341,7 @@ def insert_docstring_into_function(function_code: str, docstring: str) -> str:
363
341
  lines = function_code.splitlines(keepends=True)
364
342
 
365
343
  tree = ast.parse(function_code)
366
- if not tree.body or not isinstance(
367
- tree.body[0], ast.FunctionDef | ast.AsyncFunctionDef
368
- ):
344
+ if not tree.body or not isinstance(tree.body[0], ast.FunctionDef | ast.AsyncFunctionDef):
369
345
  print(
370
346
  "Warning: Could not parse function definition from code snippet. Returning original code.",
371
347
  file=sys.stderr,
@@ -399,15 +375,11 @@ def insert_docstring_into_function(function_code: str, docstring: str) -> str:
399
375
  if func_node.lineno - 1 < len(lines): # Ensure def line exists
400
376
  def_line = lines[func_node.lineno - 1]
401
377
  def_line_indent = def_line[: len(def_line) - len(def_line.lstrip())]
402
- body_indent = (
403
- def_line_indent + " "
404
- ) # Standard 4 spaces relative indent
378
+ body_indent = def_line_indent + " " # Standard 4 spaces relative indent
405
379
 
406
380
  # Format the new docstring lines with the calculated indentation
407
381
  new_docstring_lines_formatted = [f'{body_indent}"""\n']
408
- new_docstring_lines_formatted.extend(
409
- [f"{body_indent}{line}\n" for line in docstring.splitlines()]
410
- )
382
+ new_docstring_lines_formatted.extend([f"{body_indent}{line}\n" for line in docstring.splitlines()])
411
383
  new_docstring_lines_formatted.append(f'{body_indent}"""\n')
412
384
 
413
385
  output_lines = []
@@ -425,10 +397,7 @@ def insert_docstring_into_function(function_code: str, docstring: str) -> str:
425
397
  dummy_tree = ast.parse(dummy_code)
426
398
  dummy_body_statements = (
427
399
  dummy_tree.body[0].body
428
- if dummy_tree.body
429
- and isinstance(
430
- dummy_tree.body[0], ast.FunctionDef | ast.AsyncFunctionDef
431
- )
400
+ if dummy_tree.body and isinstance(dummy_tree.body[0], ast.FunctionDef | ast.AsyncFunctionDef)
432
401
  else []
433
402
  )
434
403
  cleaned_body_parts = []
@@ -455,13 +424,9 @@ def insert_docstring_into_function(function_code: str, docstring: str) -> str:
455
424
 
456
425
  if not is_just_string_stmt:
457
426
  stmt_start_idx = stmt_node.lineno - 1
458
- stmt_end_idx = (
459
- stmt_node.end_lineno - 1
460
- ) # Inclusive end line index
427
+ stmt_end_idx = stmt_node.end_lineno - 1 # Inclusive end line index
461
428
 
462
- cleaned_body_parts.extend(
463
- lines[stmt_start_idx : stmt_end_idx + 1]
464
- )
429
+ cleaned_body_parts.extend(lines[stmt_start_idx : stmt_end_idx + 1])
465
430
 
466
431
  if func_node.body:
467
432
  last_stmt_end_idx = func_node.body[-1].end_lineno - 1
@@ -545,9 +510,7 @@ def process_file(file_path: str, model: str = "perplexity/sonar") -> int:
545
510
  formatted_docstring = format_docstring(docstring_output)
546
511
 
547
512
  # Insert docstring into function
548
- updated_function = insert_docstring_into_function(
549
- function_code, formatted_docstring
550
- )
513
+ updated_function = insert_docstring_into_function(function_code, formatted_docstring)
551
514
 
552
515
  # Replace the function in the file content
553
516
  if updated_function != function_code: