pub(super) fn find_python_insertion_point(content: &[u8]) -> usize {
let mut pos = 0;
if content.starts_with(b"\xef\xbb\xbf") {
pos = 3;
}
let find_line_end = |start: usize| -> usize {
content[start..]
.iter()
.position(|&b| b == b'\n')
.map(|i| start + i + 1)
.unwrap_or(content.len())
};
let is_comment_or_blank = |line: &[u8]| -> bool {
let trimmed = line
.iter()
.position(|b| !b.is_ascii_whitespace())
.map(|i| &line[i..])
.unwrap_or(&[]);
trimmed.is_empty() || trimmed.starts_with(b"#")
};
let is_encoding_line = |line: &[u8]| -> bool {
let trimmed = line
.iter()
.position(|b| !b.is_ascii_whitespace())
.map(|i| &line[i..])
.unwrap_or(&[]);
if !trimmed.starts_with(b"#") {
return false;
}
let comment = &trimmed[1..];
comment
.windows(7)
.any(|w| &w[..6] == b"coding" && (w[6] == b':' || w[6] == b'='))
};
let pos_after_bom = pos;
if pos < content.len() {
let line_end = find_line_end(pos);
let line = &content[pos..line_end];
if line.starts_with(b"#!") {
pos = line_end;
}
}
let lines_to_check: usize = if pos > pos_after_bom { 1 } else { 2 };
for _ in 0..lines_to_check {
if pos >= content.len() {
break;
}
let line_end = find_line_end(pos);
let line = &content[pos..line_end];
if is_encoding_line(line) {
pos = line_end;
break;
}
}
let mut insertion_point = pos;
while pos < content.len() {
let line_end = find_line_end(pos);
let line = &content[pos..line_end];
if is_comment_or_blank(line) {
pos = line_end;
insertion_point = pos;
} else {
break;
}
}
if pos < content.len() {
let trimmed_start = content[pos..]
.iter()
.position(|b| !b.is_ascii_whitespace())
.map(|i| pos + i)
.unwrap_or(pos);
let rest = &content[trimmed_start..];
let quote = if rest.starts_with(b"\"\"\"") {
Some(b"\"\"\"".as_slice())
} else if rest.starts_with(b"'''") {
Some(b"'''".as_slice())
} else if rest.starts_with(b"r\"\"\"")
|| rest.starts_with(b"R\"\"\"")
|| rest.starts_with(b"u\"\"\"")
|| rest.starts_with(b"U\"\"\"")
{
Some(b"\"\"\"".as_slice())
} else if rest.starts_with(b"r'''")
|| rest.starts_with(b"R'''")
|| rest.starts_with(b"u'''")
|| rest.starts_with(b"U'''")
{
Some(b"'''".as_slice())
} else {
None
};
if let Some(q) = quote {
let start_offset = if rest.starts_with(b"r")
|| rest.starts_with(b"R")
|| rest.starts_with(b"u")
|| rest.starts_with(b"U")
{
4
} else {
3
};
if let Some(end_idx) = rest[start_offset..]
.windows(3)
.position(|w| w == q)
.map(|i| trimmed_start + start_offset + i + 3)
{
pos = end_idx;
if pos < content.len() && content[pos] == b'\n' {
pos += 1;
}
insertion_point = pos;
}
}
}
let mut in_future_import = false;
let mut paren_depth: usize = 0;
while pos < content.len() {
let line_end = find_line_end(pos);
let line = &content[pos..line_end];
if in_future_import {
for &b in line {
match b {
b'(' => paren_depth += 1,
b')' => paren_depth = paren_depth.saturating_sub(1),
_ => {}
}
}
let ends_with_backslash = !line.is_empty()
&& line
.iter()
.rposition(|b| !b.is_ascii_whitespace())
.map(|i| line[i] == b'\\')
.unwrap_or(false);
if paren_depth == 0 && !ends_with_backslash {
in_future_import = false;
insertion_point = line_end;
}
} else {
let trimmed_start = line
.iter()
.position(|b| !b.is_ascii_whitespace())
.unwrap_or(0);
if line[trimmed_start..].starts_with(b"from __future__") {
paren_depth = 0;
for &b in line {
match b {
b'(' => paren_depth += 1,
b')' => paren_depth = paren_depth.saturating_sub(1),
_ => {}
}
}
let ends_with_backslash = !line.is_empty()
&& line
.iter()
.rposition(|b| !b.is_ascii_whitespace())
.map(|i| line[i] == b'\\')
.unwrap_or(false);
if paren_depth > 0 || ends_with_backslash {
in_future_import = true;
} else {
insertion_point = line_end;
}
} else if !is_comment_or_blank(line) {
break;
}
}
pos = line_end;
}
insertion_point
}
#[cfg(test)]
mod tests {
#[test]
fn test_find_python_insertion_point() {
use super::find_python_insertion_point;
assert_eq!(find_python_insertion_point(b"import os\n"), 0);
let content = b"from __future__ import annotations\nimport os\n";
assert_eq!(find_python_insertion_point(content), 35);
let content =
b"from __future__ import annotations\nfrom __future__ import division\nimport os\n";
assert_eq!(find_python_insertion_point(content), 67);
let content = b"\"\"\"Docstring.\"\"\"\nfrom __future__ import annotations\nimport os\n";
assert_eq!(find_python_insertion_point(content), 52);
let content = b"\"\"\"Module docstring.\"\"\"\nimport os\n";
assert_eq!(find_python_insertion_point(content), 24);
assert_eq!(find_python_insertion_point(b""), 0);
let content = b"from __future__ import (\n annotations,\n)\nimport os\n";
assert_eq!(find_python_insertion_point(content), 44);
let content = b"from __future__ import annotations, \\\n division\nimport os\n";
assert_eq!(find_python_insertion_point(content), 51);
let content = b"\xef\xbb\xbfimport os\n";
assert_eq!(find_python_insertion_point(content), 3);
let content = b"\xef\xbb\xbffrom __future__ import annotations\nimport os\n";
assert_eq!(find_python_insertion_point(content), 38);
let content = b"#!/usr/bin/env python\nimport os\n";
assert_eq!(find_python_insertion_point(content), 22);
let content = b"#!/usr/bin/env python\n# -*- coding: utf-8 -*-\nimport os\n";
assert_eq!(find_python_insertion_point(content), 46);
let content =
b"#!/usr/bin/env python\n# -*- coding: utf-8 -*-\nfrom __future__ import annotations\nimport os\n";
assert_eq!(find_python_insertion_point(content), 81);
let content = b"# coding: utf-8\nimport os\n";
assert_eq!(find_python_insertion_point(content), 16);
let content = b"r\"\"\"Raw docstring.\"\"\"\nimport os\n";
assert_eq!(find_python_insertion_point(content), 22);
let content = b"'''Single quoted docstring.'''\nimport os\n";
assert_eq!(find_python_insertion_point(content), 31);
let content = b"\"\"\"Multi-line\ndocstring.\"\"\"\nimport os\n";
assert_eq!(find_python_insertion_point(content), 28);
}
}