1use super::{Formatter, HeaderMetadata, Import, Section};
15use std::collections::BTreeMap;
16
17#[derive(Debug, Clone)]
42pub struct PythonFormatter;
43
44impl PythonFormatter {
45 #[must_use]
47 pub const fn new() -> Self {
48 Self
49 }
50}
51
52impl Default for PythonFormatter {
53 fn default() -> Self {
54 Self::new()
55 }
56}
57
58impl Formatter for PythonFormatter {
59 fn format_header(&self, metadata: &HeaderMetadata) -> String {
60 let mut header = String::new();
61
62 header.push_str("#!/usr/bin/env python3\n");
63
64 header.push_str("# ruff: noqa: EXE001, I001\n");
65
66 if metadata.auto_generated {
67 header.push_str("# DO NOT EDIT - Auto-generated by Spikard CLI\n");
68 if let Some(schema_file) = &metadata.schema_file {
69 header.push_str(&format!("# Schema: {schema_file}\n"));
70 }
71 if let Some(version) = &metadata.generator_version {
72 header.push_str(&format!("# Generator: Spikard {version}\n"));
73 }
74 }
75
76 header.push_str("\n\"\"\"GraphQL types generated from schema.\"\"\"\n");
77
78 header
79 }
80
81 fn format_imports(&self, imports: &[Import]) -> String {
82 if imports.is_empty() {
83 return String::new();
84 }
85
86 let mut future_imports = Vec::new();
87 let mut stdlib_imports = BTreeMap::new();
88 let mut third_party_imports = BTreeMap::new();
89 let mut local_imports = BTreeMap::new();
90
91 let stdlib_modules = [
92 "abc",
93 "argparse",
94 "array",
95 "asyncio",
96 "bisect",
97 "builtins",
98 "calendar",
99 "cmath",
100 "cmd",
101 "code",
102 "codeop",
103 "collections",
104 "colorsys",
105 "compileall",
106 "concurrent",
107 "configparser",
108 "contextlib",
109 "contextvars",
110 "copy",
111 "copyreg",
112 "cprofile",
113 "csv",
114 "ctypes",
115 "curses",
116 "dataclasses",
117 "datetime",
118 "dbm",
119 "decimal",
120 "difflib",
121 "dis",
122 "doctest",
123 "email",
124 "encodings",
125 "enum",
126 "errno",
127 "faulthandler",
128 "fcntl",
129 "filecmp",
130 "fileinput",
131 "fnmatch",
132 "fractions",
133 "ftplib",
134 "functools",
135 "gc",
136 "getopt",
137 "getpass",
138 "gettext",
139 "glob",
140 "grp",
141 "gzip",
142 "hashlib",
143 "heapq",
144 "hmac",
145 "html",
146 "http",
147 "idlelib",
148 "imaplib",
149 "imghdr",
150 "imp",
151 "importlib",
152 "inspect",
153 "io",
154 "ipaddress",
155 "itertools",
156 "json",
157 "keyword",
158 "lib2to3",
159 "linecache",
160 "locale",
161 "logging",
162 "lzma",
163 "mailbox",
164 "mailcap",
165 "marshal",
166 "math",
167 "mimetypes",
168 "mmap",
169 "modulefinder",
170 "msilib",
171 "msvcrt",
172 "multiprocessing",
173 "netrc",
174 "nis",
175 "nntplib",
176 "numbers",
177 "operator",
178 "optparse",
179 "os",
180 "ossaudiodev",
181 "parser",
182 "pathlib",
183 "pdb",
184 "pickle",
185 "pickletools",
186 "pipes",
187 "pkgutil",
188 "platform",
189 "plistlib",
190 "poplib",
191 "posix",
192 "posixpath",
193 "pprint",
194 "profile",
195 "pstats",
196 "pty",
197 "pwd",
198 "py_compile",
199 "pyclbr",
200 "pydoc",
201 "queue",
202 "quopri",
203 "random",
204 "re",
205 "readline",
206 "reprlib",
207 "resource",
208 "rlcompleter",
209 "runpy",
210 "sched",
211 "secrets",
212 "select",
213 "selectors",
214 "shelve",
215 "shlex",
216 "shutil",
217 "signal",
218 "site",
219 "smtpd",
220 "smtplib",
221 "sndhdr",
222 "socket",
223 "socketserver",
224 "spwd",
225 "sqlite3",
226 "ssl",
227 "stat",
228 "statistics",
229 "string",
230 "stringprep",
231 "struct",
232 "subprocess",
233 "sunau",
234 "symbol",
235 "symtable",
236 "sys",
237 "sysconfig",
238 "syslog",
239 "tabnanny",
240 "tarfile",
241 "telnetlib",
242 "tempfile",
243 "termios",
244 "test",
245 "textwrap",
246 "threading",
247 "time",
248 "timeit",
249 "tkinter",
250 "token",
251 "tokenize",
252 "trace",
253 "traceback",
254 "tracemalloc",
255 "tty",
256 "turtle",
257 "types",
258 "typing",
259 "typing_extensions",
260 "unicodedata",
261 "unittest",
262 "urllib",
263 "uu",
264 "uuid",
265 "venv",
266 "warnings",
267 "wave",
268 "weakref",
269 "webbrowser",
270 "winreg",
271 "winsound",
272 "wsgiref",
273 "xdrlib",
274 "xml",
275 "xmlrpc",
276 "zipapp",
277 "zipfile",
278 "zipimport",
279 "zlib",
280 ];
281
282 for import in imports {
283 let module_name = import.module.split('.').next().unwrap_or(&import.module);
284
285 if module_name == "__future__" {
286 future_imports.push(import.clone());
287 } else if stdlib_modules.contains(&module_name) {
288 stdlib_imports
289 .entry(import.module.clone())
290 .or_insert_with(Vec::new)
291 .push(import.clone());
292 } else if module_name.starts_with('.') {
293 local_imports
294 .entry(import.module.clone())
295 .or_insert_with(Vec::new)
296 .push(import.clone());
297 } else {
298 third_party_imports
299 .entry(import.module.clone())
300 .or_insert_with(Vec::new)
301 .push(import.clone());
302 }
303 }
304
305 let mut output = String::new();
306
307 if !future_imports.is_empty() {
308 for import in &future_imports {
309 output.push_str(&format_python_import(import));
310 output.push('\n');
311 }
312 output.push('\n');
313 }
314
315 if !stdlib_imports.is_empty() {
316 for imports_vec in stdlib_imports.values() {
317 for import in imports_vec {
318 output.push_str(&format_python_import(import));
319 output.push('\n');
320 }
321 }
322 output.push('\n');
323 }
324
325 if !third_party_imports.is_empty() {
326 for imports_vec in third_party_imports.values() {
327 for import in imports_vec {
328 output.push_str(&format_python_import(import));
329 output.push('\n');
330 }
331 }
332 output.push('\n');
333 }
334
335 if !local_imports.is_empty() {
336 for imports_vec in local_imports.values() {
337 for import in imports_vec {
338 output.push_str(&format_python_import(import));
339 output.push('\n');
340 }
341 }
342 }
343
344 output.trim_end().to_string()
345 }
346
347 fn format_docstring(&self, content: &str) -> String {
348 let escaped = content.replace("\"\"\"", r#"\"\"\""#);
349
350 format!("\"\"\"{escaped}\"\"\"")
351 }
352
353 fn merge_sections(&self, sections: &[Section]) -> String {
354 let mut header = String::new();
355 let mut imports = String::new();
356 let mut body = String::new();
357
358 for section in sections {
359 match section {
360 Section::Header(content) => {
361 if header.is_empty() {
362 header = content.clone();
363 }
364 }
365 Section::Imports(content) => {
366 if imports.is_empty() {
367 imports = content.clone();
368 }
369 }
370 Section::Body(content) => {
371 if body.is_empty() {
372 body = content.clone();
373 }
374 }
375 }
376 }
377
378 let mut output = String::new();
379
380 if !header.is_empty() {
381 output.push_str(header.trim_end());
382 output.push_str("\n\n");
383 }
384
385 if !imports.is_empty() {
386 output.push_str(imports.trim_end());
387 output.push_str("\n\n");
388 }
389
390 if !body.is_empty() {
391 output.push_str(body.trim_end());
392 output.push('\n');
393 }
394
395 output.trim_end().to_string() + "\n"
396 }
397}
398
399fn format_python_import(import: &Import) -> String {
401 if import.items.is_empty() {
402 format!("import {}", import.module)
403 } else {
404 let items = import.items.join(", ");
405 format!("from {} import {}", import.module, items)
406 }
407}
408
409#[cfg(test)]
410mod tests {
411 use super::*;
412
413 #[test]
414 fn test_format_header_with_metadata() {
415 let formatter = PythonFormatter::new();
416 let metadata = HeaderMetadata {
417 auto_generated: true,
418 schema_file: Some("schema.graphql".to_string()),
419 generator_version: Some("0.6.2".to_string()),
420 };
421
422 let header = formatter.format_header(&metadata);
423 assert!(header.contains("#!/usr/bin/env python3"));
424 assert!(header.contains("# ruff: noqa: EXE001, I001"));
425 assert!(header.contains("# DO NOT EDIT - Auto-generated by Spikard CLI"));
426 assert!(header.contains("# Schema: schema.graphql"));
427 assert!(header.contains("# Generator: Spikard 0.6.2"));
428 assert!(header.contains("\"\"\"GraphQL types generated from schema.\"\"\""));
429 }
430
431 #[test]
432 fn test_format_header_without_metadata() {
433 let formatter = PythonFormatter::new();
434 let metadata = HeaderMetadata {
435 auto_generated: false,
436 schema_file: None,
437 generator_version: None,
438 };
439
440 let header = formatter.format_header(&metadata);
441 assert!(header.contains("#!/usr/bin/env python3"));
442 assert!(!header.contains("# DO NOT EDIT"));
443 }
444
445 #[test]
446 fn test_format_imports_empty() {
447 let formatter = PythonFormatter::new();
448 let imports = [];
449 let output = formatter.format_imports(&imports);
450 assert!(output.is_empty());
451 }
452
453 #[test]
454 fn test_format_imports_grouped_and_sorted() {
455 let formatter = PythonFormatter::new();
456 let imports = vec![
457 Import::with_items("typing", vec!["List", "Dict"]),
458 Import::with_items("__future__", vec!["annotations"]),
459 Import::new("msgspec"),
460 Import::new("graphql"),
461 ];
462
463 let output = formatter.format_imports(&imports);
464 let lines: Vec<&str> = output.lines().collect();
465
466 assert_eq!(lines[0], "from __future__ import annotations");
467 assert!(lines.contains(&"from typing import List, Dict"));
468 assert!(lines.contains(&"import graphql"));
469 assert!(lines.contains(&"import msgspec"));
470 }
471
472 #[test]
473 fn test_format_imports_simple_module() {
474 let formatter = PythonFormatter::new();
475 let imports = vec![Import::new("asyncio")];
476
477 let output = formatter.format_imports(&imports);
478 assert_eq!(output.trim(), "import asyncio");
479 }
480
481 #[test]
482 fn test_format_imports_with_items() {
483 let formatter = PythonFormatter::new();
484 let imports = vec![Import::with_items("typing", vec!["Optional", "Union"])];
485
486 let output = formatter.format_imports(&imports);
487 assert_eq!(output.trim(), "from typing import Optional, Union");
488 }
489
490 #[test]
491 fn test_format_docstring() {
492 let formatter = PythonFormatter::new();
493 let content = "This is a test docstring";
494 let output = formatter.format_docstring(content);
495 assert_eq!(output, "\"\"\"This is a test docstring\"\"\"");
496 }
497
498 #[test]
499 fn test_format_docstring_with_quotes() {
500 let formatter = PythonFormatter::new();
501 let content = r#"This says "hello""""#;
502 let output = formatter.format_docstring(content);
503 assert!(output.contains(r#"\"\"\""#));
504 }
505
506 #[test]
507 fn test_merge_sections_in_order() {
508 let formatter = PythonFormatter::new();
509 let sections = vec![
510 Section::Header("#!/usr/bin/env python3\n# Auto-gen".to_string()),
511 Section::Imports("from typing import List".to_string()),
512 Section::Body("class MyType:\n pass".to_string()),
513 ];
514
515 let output = formatter.merge_sections(§ions);
516 let lines: Vec<&str> = output.lines().collect();
517
518 assert!(lines[0].contains("#!/usr/bin/env python3"));
519 assert!(lines.iter().any(|l| l.contains("from typing import List")));
520 assert!(lines.iter().any(|l| l.contains("class MyType")));
521 }
522
523 #[test]
524 fn test_merge_sections_duplicate_headers() {
525 let formatter = PythonFormatter::new();
526 let sections = vec![
527 Section::Header("#!/usr/bin/env python3".to_string()),
528 Section::Header("#!/usr/bin/env python3".to_string()),
529 Section::Body("class MyType:\n pass".to_string()),
530 ];
531
532 let output = formatter.merge_sections(§ions);
533 let header_count = output.matches("#!/usr/bin/env python3").count();
534 assert_eq!(header_count, 1, "Should not duplicate headers");
535 }
536
537 #[test]
538 fn test_merge_sections_trailing_newline() {
539 let formatter = PythonFormatter::new();
540 let sections = vec![Section::Body("class MyType:\n pass".to_string())];
541
542 let output = formatter.merge_sections(§ions);
543 assert!(output.ends_with('\n'));
544 assert!(!output.ends_with("\n\n"));
545 }
546}