use std::collections::BTreeMap;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TagKind {
Class,
Def,
}
impl TagKind {
pub fn as_str(self) -> &'static str {
match self {
TagKind::Class => "class",
TagKind::Def => "def",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PycodeError(pub String);
impl std::fmt::Display for PycodeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
impl std::error::Error for PycodeError {}
pub fn find_tags(source: &str) -> Result<BTreeMap<String, (TagKind, u32, u32)>, PycodeError> {
let code = filter_whitespace(source);
let lines = splitlines_keepends(&code);
let tokens = tokenize(&lines)?;
DefinitionFinder::new(&tokens, &lines).parse()
}
fn filter_whitespace(code: &str) -> String {
code.replace('\x0c', " ")
}
fn splitlines_keepends(text: &str) -> Vec<String> {
let mut out = Vec::new();
let mut start = 0usize;
let mut chars = text.char_indices().peekable();
while let Some((i, c)) = chars.next() {
if is_line_boundary(c) {
let mut end = i + c.len_utf8();
if c == '\r' {
if let Some(&(j, '\n')) = chars.peek() {
chars.next();
end = j + 1;
}
}
out.push(text[start..end].to_string());
start = end;
}
}
if start < text.len() {
out.push(text[start..].to_string());
}
out
}
fn is_line_boundary(c: char) -> bool {
matches!(
c,
'\n' | '\r'
| '\x0b'
| '\x0c'
| '\x1c'
| '\x1d'
| '\x1e'
| '\u{85}'
| '\u{2028}'
| '\u{2029}'
)
}
fn is_emptyline(line: &str) -> bool {
let body = line.strip_suffix('\n').unwrap_or(line);
let rest = body.trim_start_matches(char::is_whitespace);
rest.is_empty() || rest.starts_with('#')
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Kind {
Name,
Number,
Str,
Op,
Comment,
Newline,
Nl,
Indent,
Dedent,
}
#[derive(Debug, Clone)]
struct Tok {
kind: Kind,
text: String,
start_row: u32,
end_row: u32,
}
const OPS3: &[&str] = &["**=", "//=", "<<=", ">>=", "..."];
const OPS2: &[&str] = &[
"**", "//", "<<", ">>", "<=", ">=", "==", "!=", "->", ":=", "+=", "-=", "*=", "/=", "%=", "@=",
"&=", "|=", "^=",
];
fn is_name_start(c: char) -> bool {
c.is_alphabetic() || c == '_'
}
fn is_name_continue(c: char) -> bool {
c.is_alphanumeric() || c == '_'
}
fn is_string_prefix(word: &str) -> bool {
!word.is_empty() && word.len() <= 2 && word.chars().all(|c| "rRbBuUfF".contains(c))
}
struct Tokenizer {
body: Vec<Vec<char>>,
toks: Vec<Tok>,
indents: Vec<(usize, usize)>,
brackets: Vec<(char, u32)>,
continued: bool,
line_has_tokens: bool,
row: usize,
col: usize,
}
fn tokenize(lines: &[String]) -> Result<Vec<Tok>, PycodeError> {
let body: Vec<Vec<char>> = lines
.iter()
.map(|l| line_body(l).chars().collect())
.collect();
let mut tk = Tokenizer {
body,
toks: Vec::new(),
indents: vec![(0, 0)],
brackets: Vec::new(),
continued: false,
line_has_tokens: false,
row: 0,
col: 0,
};
tk.run()?;
Ok(tk.toks)
}
fn line_body(line: &str) -> &str {
if let Some(rest) = line.strip_suffix("\r\n") {
rest
} else {
line.strip_suffix(is_line_boundary).unwrap_or(line)
}
}
impl Tokenizer {
fn push(&mut self, kind: Kind, text: &str, start_row: usize, end_row: usize) {
self.toks.push(Tok {
kind,
text: text.to_string(),
start_row: start_row as u32 + 1,
end_row: end_row as u32 + 1,
});
}
fn cur(&self) -> Option<char> {
self.body[self.row].get(self.col).copied()
}
fn at(&self, offset: usize) -> Option<char> {
self.body[self.row].get(self.col + offset).copied()
}
fn line_len(&self) -> usize {
self.body[self.row].len()
}
fn run(&mut self) -> Result<(), PycodeError> {
let mut at_line_start = true;
while self.row < self.body.len() {
if at_line_start {
at_line_start = false;
if self.brackets.is_empty() && !self.continued {
if self.start_of_logical_line()? {
self.row += 1;
self.col = 0;
at_line_start = true;
continue;
}
} else {
self.col = 0;
}
self.continued = false;
}
if self.col >= self.line_len() {
self.end_of_physical_line();
self.row += 1;
self.col = 0;
at_line_start = true;
continue;
}
self.scan_token()?;
}
self.finish()
}
fn start_of_logical_line(&mut self) -> Result<bool, PycodeError> {
let (col, altcol, first) = self.measure_indent();
self.col = first;
match self.body[self.row].get(first) {
None => {
self.push(Kind::Nl, "", self.row, self.row);
return Ok(true);
}
Some('#') => {
self.push(Kind::Comment, "", self.row, self.row);
self.push(Kind::Nl, "", self.row, self.row);
return Ok(true);
}
Some(_) => {}
}
let &(top, alttop) = self.indents.last().expect("indent stack is never empty");
if col == top {
if altcol != alttop {
return Err(self.tab_error());
}
} else if col > top {
if altcol <= alttop {
return Err(self.tab_error());
}
self.indents.push((col, altcol));
self.push(Kind::Indent, "", self.row, self.row);
} else {
while self.indents.len() > 1 && col < self.indents[self.indents.len() - 1].0 {
self.indents.pop();
self.push(Kind::Dedent, "", self.row, self.row);
}
let &(top, alttop) = self.indents.last().expect("indent stack is never empty");
if col != top {
return Err(PycodeError(
"unindent does not match any outer indentation level".to_string(),
));
}
if altcol != alttop {
return Err(self.tab_error());
}
}
Ok(false)
}
fn measure_indent(&self) -> (usize, usize, usize) {
let (mut col, mut altcol, mut i) = (0usize, 0usize, 0usize);
let line = &self.body[self.row];
while let Some(&c) = line.get(i) {
match c {
' ' => {
col += 1;
altcol += 1;
}
'\t' => {
col = (col / 8 + 1) * 8;
altcol += 1;
}
_ => break,
}
i += 1;
}
(col, altcol, i)
}
fn tab_error(&self) -> PycodeError {
PycodeError("inconsistent use of tabs and spaces in indentation".to_string())
}
fn end_of_physical_line(&mut self) {
if self.continued {
return;
}
if !self.brackets.is_empty() {
self.push(Kind::Nl, "", self.row, self.row);
} else if self.line_has_tokens {
self.push(Kind::Newline, "", self.row, self.row);
self.line_has_tokens = false;
} else {
self.push(Kind::Nl, "", self.row, self.row);
}
}
fn scan_token(&mut self) -> Result<(), PycodeError> {
let c = self.cur().expect("caller checked the column");
if c == ' ' || c == '\t' || c == '\r' {
self.col += 1;
return Ok(());
}
if c == '#' {
self.push(Kind::Comment, "", self.row, self.row);
self.col = self.line_len();
return Ok(());
}
if c == '\\' && self.col + 1 >= self.line_len() {
self.continued = true;
self.col = self.line_len();
return Ok(());
}
if is_name_start(c) {
let start = self.col;
while self.cur().is_some_and(is_name_continue) {
self.col += 1;
}
let word: String = self.body[self.row][start..self.col].iter().collect();
if matches!(self.cur(), Some('\'' | '"')) && is_string_prefix(&word) {
return self.lex_string();
}
self.line_has_tokens = true;
self.push(Kind::Name, &word, self.row, self.row);
return Ok(());
}
if c.is_ascii_digit() || (c == '.' && self.at(1).is_some_and(|d| d.is_ascii_digit())) {
self.lex_number();
return Ok(());
}
if c == '\'' || c == '"' {
return self.lex_string();
}
self.lex_op();
Ok(())
}
fn lex_number(&mut self) {
let line = &self.body[self.row];
let start = self.col;
let radix_prefixed = line.get(start) == Some(&'0')
&& matches!(line.get(start + 1), Some('x' | 'X' | 'b' | 'B' | 'o' | 'O'));
let mut prev = line[start];
self.col += 1;
while let Some(d) = self.cur() {
let continues = d.is_ascii_alphanumeric()
|| d == '_'
|| d == '.'
|| ((d == '+' || d == '-') && matches!(prev, 'e' | 'E') && !radix_prefixed);
if !continues {
break;
}
prev = d;
self.col += 1;
}
self.line_has_tokens = true;
self.push(Kind::Number, "", self.row, self.row);
}
fn lex_string(&mut self) -> Result<(), PycodeError> {
let start_row = self.row;
let quote = self.cur().expect("caller peeked the quote");
self.col += 1;
let triple = self.cur() == Some(quote) && self.at(1) == Some(quote);
if triple {
self.col += 2;
}
loop {
let Some(c) = self.cur() else {
if !triple {
return Err(PycodeError(format!(
"unterminated string literal (detected at line {})",
start_row + 1
)));
}
if self.row + 1 >= self.body.len() {
return Err(PycodeError(format!(
"unterminated triple-quoted string literal (detected at line {})",
self.body.len()
)));
}
self.row += 1;
self.col = 0;
continue;
};
if c == '\\' {
self.col += 1;
if self.cur().is_none() {
if self.row + 1 >= self.body.len() {
return Err(PycodeError(format!(
"unterminated string literal (detected at line {})",
start_row + 1
)));
}
self.row += 1;
self.col = 0;
} else {
self.col += 1;
}
continue;
}
if c == quote {
if !triple {
self.col += 1;
break;
}
if self.at(1) == Some(quote) && self.at(2) == Some(quote) {
self.col += 3;
break;
}
}
self.col += 1;
}
self.line_has_tokens = true;
let end_row = self.row;
self.push(Kind::Str, "", start_row, end_row);
Ok(())
}
fn lex_op(&mut self) {
let rest: String = self.body[self.row][self.col..].iter().collect();
let text = OPS3
.iter()
.chain(OPS2.iter())
.find(|cand| rest.starts_with(**cand))
.map(|cand| (*cand).to_string())
.unwrap_or_else(|| rest.chars().next().into_iter().collect());
self.col += text.chars().count();
match text.as_str() {
"(" | "[" | "{" => self
.brackets
.push((text.chars().next().expect("one char"), self.row as u32 + 1)),
")" | "]" | "}" => {
self.brackets.pop();
}
_ => {}
}
self.line_has_tokens = true;
self.push(Kind::Op, &text, self.row, self.row);
}
fn finish(&mut self) -> Result<(), PycodeError> {
if let Some(&(opener, row)) = self.brackets.first() {
return Err(PycodeError(format!(
"'{opener}' was never closed (opened at line {row})"
)));
}
let eof_row = self.body.len();
while self.indents.len() > 1 {
self.indents.pop();
self.push(Kind::Dedent, "", eof_row, eof_row);
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Block {
Other,
Class,
Def,
}
impl Block {
fn of(kind: TagKind) -> Self {
match kind {
TagKind::Class => Block::Class,
TagKind::Def => Block::Def,
}
}
}
#[derive(Debug, Clone, Copy)]
enum Cond {
Colon,
Indent,
}
struct DefinitionFinder<'a> {
toks: &'a [Tok],
lines: &'a [String],
next: usize,
current: Option<usize>,
previous: Option<usize>,
decorator: Option<u32>,
context: Vec<String>,
indents: Vec<(Block, String, u32)>,
definitions: BTreeMap<String, (TagKind, u32, u32)>,
}
impl<'a> DefinitionFinder<'a> {
fn new(toks: &'a [Tok], lines: &'a [String]) -> Self {
Self {
toks,
lines,
next: 0,
current: None,
previous: None,
decorator: None,
context: Vec::new(),
indents: Vec::new(),
definitions: BTreeMap::new(),
}
}
fn tok(&self, i: usize) -> &'a Tok {
&self.toks[i]
}
fn fetch(&mut self) -> Option<usize> {
self.previous = self.current;
self.current = (self.next < self.toks.len()).then(|| {
let i = self.next;
self.next += 1;
i
});
self.current
}
fn parse(mut self) -> Result<BTreeMap<String, (TagKind, u32, u32)>, PycodeError> {
while let Some(i) = self.fetch() {
let t = self.tok(i);
match t.kind {
Kind::Comment => {}
Kind::Op if t.text == "@" => {
if self.decorator.is_none() && self.previous_ends_a_line() {
self.decorator = Some(t.start_row);
}
}
Kind::Name if t.text == "class" => self.parse_definition(TagKind::Class)?,
Kind::Name if t.text == "def" => self.parse_definition(TagKind::Def)?,
Kind::Indent => self.indents.push((Block::Other, String::new(), 0)),
Kind::Dedent => self.finalize_block()?,
_ => {}
}
}
Ok(self.definitions)
}
fn previous_ends_a_line(&self) -> bool {
match self.previous {
None => true,
Some(i) => matches!(
self.tok(i).kind,
Kind::Newline | Kind::Nl | Kind::Indent | Kind::Dedent
),
}
}
fn add_definition(&mut self, name: String, entry: (TagKind, u32, u32)) {
if entry.0 == TagKind::Def && self.indents.last().map(|e| e.0) == Some(Block::Def) {
return;
}
self.definitions.insert(name, entry);
}
fn parse_definition(&mut self, typ: TagKind) -> Result<(), PycodeError> {
let Some(ni) = self.fetch() else {
return Err(PycodeError(
"unexpected end of file after a definition keyword".to_string(),
));
};
let name = self.tok(ni);
let (name_text, name_end) = (name.text.clone(), name.end_row);
let start_pos = self.decorator.take().unwrap_or(name.start_row);
self.context.push(name_text);
let funcname = self.context.join(".");
self.fetch_until(Cond::Colon);
let Some(ti) = self.fetch() else {
return Err(PycodeError(
"unexpected end of file inside a definition header".to_string(),
));
};
if matches!(self.tok(ti).kind, Kind::Comment | Kind::Newline) {
self.fetch_until(Cond::Indent);
self.indents.push((Block::of(typ), funcname, start_pos));
} else {
self.add_definition(funcname, (typ, start_pos, name_end));
let _ = self.context.pop();
}
Ok(())
}
fn finalize_block(&mut self) -> Result<(), PycodeError> {
let Some((block, funcname, start_pos)) = self.indents.pop() else {
return Err(PycodeError(
"unbalanced indentation while scanning definitions".to_string(),
));
};
if block == Block::Other {
return Ok(());
}
let dedent_row = self.current.map_or(0, |i| self.tok(i).end_row);
let mut end_pos = dedent_row.saturating_sub(1);
while end_pos >= 1
&& (end_pos as usize) <= self.lines.len()
&& is_emptyline(&self.lines[end_pos as usize - 1])
{
end_pos -= 1;
}
let typ = if block == Block::Class {
TagKind::Class
} else {
TagKind::Def
};
self.add_definition(funcname, (typ, start_pos, end_pos));
let _ = self.context.pop();
Ok(())
}
fn fetch_until(&mut self, cond: Cond) {
let mut closers: Vec<&'static str> = Vec::new();
while let Some(i) = self.fetch() {
let t = self.tok(i);
let hit = match closers.last() {
Some(closer) => t.kind == Kind::Op && t.text == *closer,
None => match cond {
Cond::Colon => t.kind == Kind::Op && t.text == ":",
Cond::Indent => t.kind == Kind::Indent,
},
};
if hit {
if closers.pop().is_none() {
return;
}
continue;
}
if t.kind == Kind::Op {
match t.text.as_str() {
"(" => closers.push(")"),
"{" => closers.push("}"),
"[" => closers.push("]"),
_ => {}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn tags(src: &str) -> Vec<(String, &'static str, u32, u32)> {
find_tags(src)
.expect("find_tags must succeed")
.into_iter()
.map(|(name, (kind, start, end))| (name, kind.as_str(), start, end))
.collect()
}
fn t(name: &str, kind: &'static str, start: u32, end: u32) -> (String, &'static str, u32, u32) {
(name.to_string(), kind, start, end)
}
const FIXTURE: &str = include_str!("../../tests/fixtures/literalinclude/example.py");
#[test]
fn fixture_module_tags_match_the_probe() {
assert_eq!(
tags(FIXTURE),
vec![
t("Foo", "class", 11, 17),
t("Foo.method", "def", 16, 17),
t("tail", "def", 20, 21),
t("top", "def", 6, 8),
]
);
}
#[test]
fn a_decorator_line_becomes_the_start_for_defs_and_classes() {
let src = "import functools\n\n\n@functools.cache\ndef cached(x):\n return x\n\n\n\
@staticmethod\nclass Weird:\n pass\n";
assert_eq!(
tags(src),
vec![t("Weird", "class", 9, 11), t("cached", "def", 4, 6)]
);
}
#[test]
fn argumented_stacked_and_detached_decorators_all_start_at_the_first_at() {
assert_eq!(
tags("@decorator(\n \"a\",\n \"b\",\n)\ndef f():\n return 1\n"),
vec![t("f", "def", 1, 6)]
);
assert_eq!(
tags("@one\n@two(3)\n@three\ndef f():\n pass\n"),
vec![t("f", "def", 1, 5)]
);
assert_eq!(
tags("@deco\n\n# a comment\ndef f():\n pass\n"),
vec![t("f", "def", 1, 5)]
);
}
#[test]
fn async_def_needs_no_special_case() {
assert_eq!(
tags("@deco\nasync def f(a, b):\n await g()\n"),
vec![t("f", "def", 1, 3)]
);
assert_eq!(tags("async def f(): return 1\n"), vec![t("f", "def", 1, 1)]);
assert_eq!(
tags(
"async def outer():\n async def inner():\n pass\n\n\
\x20 class Inner:\n async def m(self):\n pass\n"
),
vec![
t("outer", "def", 1, 7),
t("outer.Inner", "class", 5, 7),
t("outer.Inner.m", "def", 6, 7),
]
);
}
#[test]
fn matrix_multiplication_is_not_a_decorator() {
assert_eq!(
tags("a = b\nc = a @ b\ndef f():\n pass\n"),
vec![t("f", "def", 3, 4)]
);
}
#[test]
fn an_at_after_an_in_bracket_newline_is_mistaken_for_a_decorator() {
assert_eq!(
tags("x = (a\n@ b)\ndef f():\n pass\n"),
vec![t("f", "def", 2, 4)]
);
assert_eq!(
tags("m = (\n a\n @ b\n)\n\n\ndef f():\n pass\n"),
vec![t("f", "def", 3, 8)]
);
}
#[test]
fn one_liners_end_at_their_header_line() {
assert_eq!(
tags("def f(): return 1\nclass C: pass\ndef g():\n pass\n"),
vec![
t("C", "class", 2, 2),
t("f", "def", 1, 1),
t("g", "def", 3, 4),
]
);
}
#[test]
fn a_one_liner_with_a_wrapped_signature_ends_at_the_name_line() {
assert_eq!(
tags("def f(a,\n b): return a + b\nx = 1\n"),
vec![t("f", "def", 1, 1)]
);
}
#[test]
fn a_backslash_continued_body_is_still_a_one_liner() {
assert_eq!(
tags("def f(a): \\\n return a\n\n\nx = 1\n"),
vec![t("f", "def", 1, 1)]
);
}
#[test]
fn semicolon_bodies_stay_one_liners() {
assert_eq!(
tags("def f(): x = 1; return x\ndef g(): pass\n"),
vec![t("f", "def", 1, 1), t("g", "def", 2, 2)]
);
}
#[test]
fn a_def_directly_inside_a_def_is_dropped() {
assert_eq!(
tags("def outer():\n def inner():\n pass\n return inner\n"),
vec![t("outer", "def", 1, 4)]
);
assert_eq!(
tags("def outer():\n def inner(): pass\n return inner\n"),
vec![t("outer", "def", 1, 3)]
);
}
#[test]
fn an_intervening_if_block_defeats_the_nested_def_suppression() {
assert_eq!(
tags("def outer():\n if True:\n def inner():\n pass\n return 1\n"),
vec![t("outer", "def", 1, 5), t("outer.inner", "def", 3, 4)]
);
assert_eq!(
tags("def outer():\n if True:\n def inner(): pass\n return 1\n"),
vec![t("outer", "def", 1, 4), t("outer.inner", "def", 3, 3)]
);
}
#[test]
fn a_def_inside_a_plain_block_keeps_its_bare_name() {
assert_eq!(
tags("if True:\n def g():\n pass\nelse:\n def h():\n pass\n"),
vec![t("g", "def", 2, 3), t("h", "def", 5, 6)]
);
assert_eq!(
tags(
"try:\n def a():\n pass\nexcept Exception:\n pass\n\n\
with open('x') as f:\n def b():\n pass\n\n\
for i in range(3):\n def c():\n pass\n"
),
vec![
t("a", "def", 2, 3),
t("b", "def", 8, 9),
t("c", "def", 12, 13),
]
);
}
#[test]
fn nested_classes_dot_their_names_and_trim_their_own_tails() {
assert_eq!(
tags(
"class Outer:\n class Inner:\n def method(self):\n pass\n\n\
\x20 attr = 1\n x = 2\n"
),
vec![
t("Outer", "class", 1, 7),
t("Outer.Inner", "class", 2, 6),
t("Outer.Inner.method", "def", 3, 4),
]
);
}
#[test]
fn a_class_inside_a_def_survives_with_its_methods() {
assert_eq!(
tags("def outer():\n class Inner:\n def m(self):\n pass\n return Inner\n"),
vec![
t("outer", "def", 1, 5),
t("outer.Inner", "class", 2, 4),
t("outer.Inner.m", "def", 3, 4),
]
);
}
#[test]
fn eof_dedents_close_every_open_block_at_the_same_line() {
assert_eq!(
tags(
"class A:\n class B:\n class C:\n def m(self):\n\
\x20 def inner():\n pass\n return inner\n"
),
vec![
t("A", "class", 1, 7),
t("A.B", "class", 2, 7),
t("A.B.C", "class", 3, 7),
t("A.B.C.m", "def", 4, 7),
]
);
}
#[test]
fn definitions_inside_strings_are_not_tags() {
let src = "DOC = \"\"\"\ndef fake():\n pass\n\nclass Fake:\n pass\n\"\"\"\n\n\n\
def real():\n \"\"\"Doc with def inside.\n\n class AlsoFake:\n pass\n \"\"\"\n return 1\n";
assert_eq!(tags(src), vec![t("real", "def", 10, 16)]);
}
#[test]
fn prefixed_string_literals_are_skipped_whole() {
let src = "s = f\"\"\"def nope():\n pass\"\"\"\nr = r\"\"\"class Nope: pass\"\"\"\n\
b = b\"def nope2(): pass\"\n\n\ndef real():\n pass\n";
assert_eq!(tags(src), vec![t("real", "def", 7, 8)]);
}
#[test]
fn single_quoted_strings_in_a_body_are_skipped() {
assert_eq!(
tags("def f():\n s = \"def nope(): pass\"\n t = 'class Nope: pass'\n return s + t\n"),
vec![t("f", "def", 1, 4)]
);
}
#[test]
fn comment_lines_make_no_tags_and_are_trimmed_from_block_ends() {
assert_eq!(
tags("# def commented():\n# pass\ndef real():\n pass\n# def trailing():\n"),
vec![t("real", "def", 3, 4)]
);
assert_eq!(
tags(
"def f():\n pass\n # trailing comment\n # another\n\n\
\x20 # after a blank\n\n\ndef g():\n pass\n"
),
vec![t("f", "def", 1, 2), t("g", "def", 9, 10)]
);
}
#[test]
fn a_column_zero_comment_does_not_close_a_block() {
assert_eq!(
tags(
"class C:\n def m(self):\n pass\n# comment at col 0\n\
\x20 def n(self):\n pass\n"
),
vec![
t("C", "class", 1, 6),
t("C.m", "def", 2, 3),
t("C.n", "def", 5, 6),
]
);
}
#[test]
fn continuation_line_signatures_span_to_the_block_end() {
assert_eq!(
tags(
"def f(\n a,\n b,\n):\n return a\n\n\nclass C(\n Base,\n):\n pass\n"
),
vec![t("C", "class", 8, 11), t("f", "def", 1, 5)]
);
}
#[test]
fn only_the_top_level_colon_closes_a_definition_header() {
assert_eq!(
tags(
"def f(a: int = 1, b: dict[str, int] = {}) -> dict[str, int]:\n return b\n\n\n\
def g(h=lambda x: x):\n return h\n"
),
vec![t("f", "def", 1, 2), t("g", "def", 5, 6)]
);
assert_eq!(
tags("def f(cb={'k': lambda x: x}):\n pass\n"),
vec![t("f", "def", 1, 2)]
);
assert_eq!(
tags("def f():\n d = {'a': 1}\n if (n := len(d)) > 0:\n return n\n"),
vec![t("f", "def", 1, 4)]
);
assert_eq!(
tags("def f(a={1: {2: 3}}, b=[1, 2], *, c: \"x\" = (1,)):\n pass\n"),
vec![t("f", "def", 1, 2)]
);
assert_eq!(
tags("def f(x=f\"{1:>10}\"):\n return x\n"),
vec![t("f", "def", 1, 2)]
);
}
#[test]
fn pep695_type_parameter_lists_are_skipped() {
assert_eq!(
tags("def f[T](x: T) -> T:\n return x\n\n\nclass C[T]:\n pass\n"),
vec![t("C", "class", 5, 6), t("f", "def", 1, 2)]
);
}
#[test]
fn keyword_prefixed_names_are_not_keywords() {
assert_eq!(
tags("class_ = 1\ndefine = 2\ndefx = 3\n\n\ndef deffo():\n pass\n"),
vec![t("deffo", "def", 6, 7)]
);
}
#[test]
fn a_redefined_name_keeps_the_last_definition() {
assert_eq!(
tags("def f():\n pass\n\n\ndef f():\n return 2\n"),
vec![t("f", "def", 5, 6)]
);
assert_eq!(
tags(
"class C:\n @property\n def p(self):\n return 1\n\n\
\x20 @p.setter\n def p(self, v):\n self._p = v\n"
),
vec![t("C", "class", 1, 8), t("C.p", "def", 6, 8)]
);
}
#[test]
fn match_statements_are_plain_blocks() {
assert_eq!(
tags("def f(x):\n match x:\n case 1:\n pass\n case _:\n pass\n"),
vec![t("f", "def", 1, 6)]
);
}
#[test]
fn tab_indentation_works_like_cpythons() {
assert_eq!(
tags("def f():\n\treturn 1\n\n\ndef g():\n\tpass\n"),
vec![t("f", "def", 1, 2), t("g", "def", 5, 6)]
);
assert_eq!(
tags("class C:\n\tdef m(self):\n\t\tpass\n\tx = 1\n"),
vec![t("C", "class", 1, 4), t("C.m", "def", 2, 3)]
);
assert_eq!(
tags("def f():\n\tpass\n # eight spaces comment\n"),
vec![t("f", "def", 1, 2)]
);
}
#[test]
fn inconsistent_tabs_and_spaces_err() {
let err =
find_tags("class C:\n\tdef m(self):\n\t\tpass\n x = 1\n").expect_err("TabError");
assert_eq!(
err.to_string(),
"inconsistent use of tabs and spaces in indentation"
);
}
#[test]
fn crlf_line_endings_number_lines_the_same() {
assert_eq!(
tags("def f():\r\n return 1\r\n\r\n\r\ndef g():\r\n pass\r\n"),
vec![t("f", "def", 1, 2), t("g", "def", 5, 6)]
);
}
#[test]
fn missing_and_surplus_trailing_newlines_both_land_on_the_last_code_line() {
assert_eq!(tags("def f():\n return 1"), vec![t("f", "def", 1, 2)]);
assert_eq!(tags("def f(): return 1"), vec![t("f", "def", 1, 1)]);
assert_eq!(
tags("def f():\n pass\n\n\n\n"),
vec![t("f", "def", 1, 2)]
);
}
#[test]
fn a_form_feed_becomes_a_blank_line_not_a_line_break() {
assert_eq!(
tags("def f():\n pass\n\x0c\n\ndef g():\n pass\n"),
vec![t("f", "def", 1, 2), t("g", "def", 5, 6)]
);
}
#[test]
fn files_without_definitions_yield_no_tags() {
assert!(tags("x = 1\ny = 2\nprint(x + y)\n").is_empty());
assert!(tags("").is_empty());
assert!(tags("# hello\n# world\n").is_empty());
assert!(tags("\n\n\n").is_empty());
}
#[test]
fn scanner_level_failures_err_with_their_own_detail() {
assert_eq!(
find_tags("x = \"\"\"abc\ndef f():\n pass\n")
.expect_err("unterminated triple")
.to_string(),
"unterminated triple-quoted string literal (detected at line 3)"
);
assert_eq!(
find_tags("x = 'abc\ndef f():\n pass\n")
.expect_err("unterminated string")
.to_string(),
"unterminated string literal (detected at line 1)"
);
assert_eq!(
find_tags("x = [1,\ndef f():\n pass\n")
.expect_err("unclosed bracket")
.to_string(),
"'[' was never closed (opened at line 1)"
);
assert_eq!(
find_tags("def f():\n pass\n x = 1\n")
.expect_err("bad dedent")
.to_string(),
"unindent does not match any outer indentation level"
);
}
#[test]
fn a_tokenizable_but_unparsable_file_still_yields_tags_here() {
assert_eq!(
tags("def f():\n return 1\n\n\nx = = 1\n"),
vec![t("f", "def", 1, 2)]
);
}
proptest::proptest! {
#[test]
fn arbitrary_source_never_panics(src in "(?s).{0,400}") {
if let Ok(map) = find_tags(&src) {
for (_, (_, start, end)) in map {
proptest::prop_assert!(start >= 1);
proptest::prop_assert!(start <= end);
}
}
}
#[test]
fn python_shaped_fragments_never_panic(
parts in proptest::collection::vec(
proptest::sample::select(vec![
"def f():", "class C:", "@deco", "async def g(): pass", " pass",
"\tpass", "x = (", ")", "\"\"\"", "'", "#", "\\", " # c", "",
"def h(a,", "):", "if True:", " deep", "\x0c", "\r",
]),
0..24,
)
) {
let src = parts.join("\n");
if let Ok(map) = find_tags(&src) {
for (_, (_, start, end)) in map {
proptest::prop_assert!(start >= 1);
proptest::prop_assert!(start <= end);
}
}
}
}
}