|
| 1 | +//! reverse of `crate::transforms::reified_generic`: |
| 2 | +//! `@generic # basedpython: reified\ndef f[T]: …` → `def f[T]: …` |
| 3 | +//! |
| 4 | +//! the forward transform tags the decorator line it synthesizes with the |
| 5 | +//! [`REIFIED_MARKER`](crate::transforms::reified_generic::REIFIED_MARKER) |
| 6 | +//! comment — provenance that this `@generic` came from reification, not from a |
| 7 | +//! user's own decorator. only a `@generic` carrying that marker is unwrapped; |
| 8 | +//! a hand-written `@generic` (no marker) is left untouched. the `generic` |
| 9 | +//! polyfill class and its `dataclasses` / `types` imports are dead once the |
| 10 | +//! wrapper is removed, and `prune_imports` drops them on the way out. |
| 11 | +
|
| 12 | +use ruff_diagnostics::{Edit, Fix}; |
| 13 | +use ruff_python_ast::visitor::{Visitor, walk_stmt}; |
| 14 | +use ruff_python_ast::{Decorator, Expr, Stmt, StmtFunctionDef}; |
| 15 | +use ruff_text_size::{Ranged, TextRange, TextSize}; |
| 16 | + |
| 17 | +use crate::transforms::reified_generic::REIFIED_MARKER; |
| 18 | +use crate::transforms::source_util::line_start; |
| 19 | + |
| 20 | +pub(crate) struct ReifiedGenericReverse<'src> { |
| 21 | + source: &'src str, |
| 22 | + pub(crate) edits: Vec<Fix>, |
| 23 | +} |
| 24 | + |
| 25 | +impl<'src> ReifiedGenericReverse<'src> { |
| 26 | + pub(crate) fn new(source: &'src str) -> Self { |
| 27 | + Self { |
| 28 | + source, |
| 29 | + edits: Vec::new(), |
| 30 | + } |
| 31 | + } |
| 32 | + |
| 33 | + /// the source position the `def`/`async` header begins at — either the next |
| 34 | + /// decorator's start, or the header keyword following this decorator |
| 35 | + fn next_header_start(&self, decorators: &[Decorator], idx: usize) -> Option<TextSize> { |
| 36 | + if let Some(next) = decorators.get(idx + 1) { |
| 37 | + return Some(next.range().start()); |
| 38 | + } |
| 39 | + let after_dec = usize::from(decorators[idx].range().end()); |
| 40 | + let rest = &self.source[after_dec..]; |
| 41 | + // skip the marker comment to the newline, then to the header keyword |
| 42 | + let offset = rest.find("def")?; |
| 43 | + Some(TextSize::from(u32::try_from(after_dec + offset).ok()?)) |
| 44 | + } |
| 45 | + |
| 46 | + /// whether the decorator is a bare `@generic` whose line carries the |
| 47 | + /// reified-provenance marker comment. the marker lives in trivia (not the |
| 48 | + /// AST), so it is matched against the source slice from the decorator's end |
| 49 | + /// to the end of its physical line |
| 50 | + fn is_marked_generic(&self, decorator: &Decorator) -> bool { |
| 51 | + if !matches!(&decorator.expression, Expr::Name(n) if n.id.as_str() == "generic") { |
| 52 | + return false; |
| 53 | + } |
| 54 | + let after = usize::from(decorator.range().end()); |
| 55 | + let rest = &self.source[after..]; |
| 56 | + let line = rest.split('\n').next().unwrap_or(rest); |
| 57 | + line.contains(REIFIED_MARKER.trim_start()) |
| 58 | + } |
| 59 | + |
| 60 | + fn unwrap_function(&mut self, function: &StmtFunctionDef) { |
| 61 | + let decorators = &function.decorator_list; |
| 62 | + for (idx, decorator) in decorators.iter().enumerate() { |
| 63 | + if !self.is_marked_generic(decorator) { |
| 64 | + continue; |
| 65 | + } |
| 66 | + // delete from the decorator's `@` through to the next header start |
| 67 | + // (the following decorator, or the `def` keyword). this removes the |
| 68 | + // marker comment and the line break with it |
| 69 | + let Some(next_start) = self.next_header_start(decorators, idx) else { |
| 70 | + continue; |
| 71 | + }; |
| 72 | + let start = line_start(self.source, decorator.range().start()); |
| 73 | + self.edits |
| 74 | + .push(Fix::safe_edit(Edit::range_deletion(TextRange::new( |
| 75 | + start, next_start, |
| 76 | + )))); |
| 77 | + } |
| 78 | + } |
| 79 | +} |
| 80 | + |
| 81 | +impl<'ast> Visitor<'ast> for ReifiedGenericReverse<'_> { |
| 82 | + fn visit_stmt(&mut self, stmt: &'ast Stmt) { |
| 83 | + if let Stmt::FunctionDef(function) = stmt { |
| 84 | + self.unwrap_function(function); |
| 85 | + } |
| 86 | + walk_stmt(self, stmt); |
| 87 | + } |
| 88 | +} |
| 89 | + |
| 90 | +#[cfg(test)] |
| 91 | +mod tests { |
| 92 | + use crate::transforms::reified_generic::REIFIED_MARKER; |
| 93 | + use crate::{Config, reverse_transpile}; |
| 94 | + use ruff_python_ast::PythonVersion; |
| 95 | + |
| 96 | + fn rev(source: &str) -> String { |
| 97 | + reverse_transpile(source, &Config::test_default()).unwrap() |
| 98 | + } |
| 99 | + |
| 100 | + #[test] |
| 101 | + fn marked_generic_is_unwrapped() { |
| 102 | + let src = format!("@generic{REIFIED_MARKER}\ndef f[T]():\n print(T)\n"); |
| 103 | + let out = rev(&src); |
| 104 | + assert!( |
| 105 | + !out.contains("@generic"), |
| 106 | + "wrapper should be removed: {out}" |
| 107 | + ); |
| 108 | + assert!(out.contains("def f[T]():"), "def should remain: {out}"); |
| 109 | + } |
| 110 | + |
| 111 | + #[test] |
| 112 | + fn handwritten_generic_is_preserved() { |
| 113 | + // no marker — a user's own `@generic` decorator stays put |
| 114 | + let src = "@generic\ndef f(x):\n return x\n"; |
| 115 | + let out = rev(src); |
| 116 | + assert!( |
| 117 | + out.contains("@generic"), |
| 118 | + "hand-written decorator must be preserved: {out}" |
| 119 | + ); |
| 120 | + } |
| 121 | + |
| 122 | + #[test] |
| 123 | + fn round_trip_rewraps() { |
| 124 | + // reverse then forward reproduces the wrapper |
| 125 | + let src = format!("@generic{REIFIED_MARKER}\ndef f[T]():\n print(T)\n"); |
| 126 | + let bare = rev(&src); |
| 127 | + let config = Config { |
| 128 | + min_version: PythonVersion::PY312, |
| 129 | + ..Config::test_default() |
| 130 | + }; |
| 131 | + let forward = crate::transpile(&bare, &config).unwrap(); |
| 132 | + assert!( |
| 133 | + forward.contains("@generic # basedpython: reified"), |
| 134 | + "forward should re-wrap the bare reified def: {forward}" |
| 135 | + ); |
| 136 | + } |
| 137 | +} |
0 commit comments