1use crate::{Diagnostic, DiagnosticCode as Code, Limits, Span, Token, TokenKind};
4
5const KEYWORDS: &[&str] = &[
6 "False", "None", "True", "and", "as", "assert", "async", "await", "break", "class", "continue",
7 "def", "del", "elif", "else", "except", "finally", "for", "from", "global", "if", "import",
8 "in", "is", "lambda", "nonlocal", "not", "or", "pass", "raise", "return", "try", "while",
9 "with", "yield",
10];
11const OPS: &[&str] = &[
12 "**=", ">>=", "<<=", "//=", "...", "->", ":=", "==", "!=", "<=", ">=", "<<", ">>", "**", "//",
13 "+=", "-=", "*=", "/=", "%=", "@=", "&=", "|=", "^=", "=>", "+", "-", "*", "/", "%", "@", "&",
14 "|", "^", "~", ":", ",", ";", ".", "=", "(", ")", "[", "]", "{", "}", "<", ">",
15];
16
17pub fn tokenize(source: &str) -> Result<Vec<Token>, Diagnostic> {
19 tokenize_with_limits(source, Limits::default())
20}
21
22pub fn tokenize_with_limits(source: &str, limits: Limits) -> Result<Vec<Token>, Diagnostic> {
24 Lexer::new(source, limits).run()
25}
26
27struct Lexer<'a> {
28 source: &'a str,
29 limits: Limits,
30 pos: usize,
31 line: usize,
32 column: usize,
33 at_line_start: bool,
34 bracket_depth: usize,
35 indents: Vec<(usize, String)>,
36 out: Vec<Token>,
37}
38
39impl<'a> Lexer<'a> {
40 fn new(source: &'a str, limits: Limits) -> Self {
41 Self {
42 source,
43 limits,
44 pos: 0,
45 line: 1,
46 column: 0,
47 at_line_start: true,
48 bracket_depth: 0,
49 indents: vec![(0, String::new())],
50 out: Vec::new(),
51 }
52 }
53
54 fn run(mut self) -> Result<Vec<Token>, Diagnostic> {
55 if self.source.len() > self.limits.max_bytes {
56 return Err(self.error(Code::ResourceLimit, 0, "source byte limit exceeded"));
57 }
58 while self.pos < self.source.len() {
59 if self.line > self.limits.max_lines {
60 return Err(self.error(
61 Code::ResourceLimit,
62 self.pos,
63 "physical line limit exceeded",
64 ));
65 }
66 if self.at_line_start && self.bracket_depth == 0 {
67 self.layout()?;
68 }
69 if self.pos >= self.source.len() {
70 break;
71 }
72 let start = self.pos;
73 let ch = self.peek().expect("position is in source");
74 if ch == '\n' || ch == '\r' {
75 self.newline()?;
76 continue;
77 }
78 if matches!(ch, ' ' | '\t' | '\u{0c}') {
79 self.take_while(|c| matches!(c, ' ' | '\t' | '\u{0c}'));
80 self.emit(TokenKind::Trivia, start)?;
81 continue;
82 }
83 if ch == '#' {
84 self.take_while(|c| !matches!(c, '\r' | '\n'));
85 self.emit(TokenKind::Trivia, start)?;
86 continue;
87 }
88 if ch == '\\'
89 && self
90 .rest()
91 .get(1..)
92 .is_some_and(|s| s.starts_with('\n') || s.starts_with("\r\n"))
93 {
94 self.bump();
95 if self.peek() == Some('\r') {
96 self.bump();
97 }
98 self.bump();
99 self.emit(TokenKind::Trivia, start)?;
100 continue;
101 }
102 if is_name_start(ch) {
103 self.name_or_string(start)?;
104 continue;
105 }
106 if ch.is_ascii_digit()
107 || (ch == '.'
108 && self
109 .rest()
110 .get(1..)
111 .is_some_and(|s| s.starts_with(|c: char| c.is_ascii_digit())))
112 {
113 self.number(start)?;
114 continue;
115 }
116 if matches!(ch, '\'' | '"') {
117 self.string(start, "", TokenKind::String)?;
118 continue;
119 }
120 if let Some(op) = OPS.iter().find(|op| self.rest().starts_with(**op)) {
121 let op = *op;
122 for _ in op.chars() {
123 self.bump();
124 }
125 match op {
126 "(" | "[" | "{" => {
127 self.bracket_depth += 1;
128 if self.bracket_depth > self.limits.max_nesting {
129 return Err(self.error(
130 Code::ResourceLimit,
131 start,
132 "delimiter nesting limit exceeded",
133 ));
134 }
135 }
136 ")" | "]" | "}" => self.bracket_depth = self.bracket_depth.saturating_sub(1),
137 _ => {}
138 }
139 self.emit(TokenKind::Operator, start)?;
140 continue;
141 }
142 self.bump();
143 return Err(self.error(
144 Code::InvalidCharacter,
145 start,
146 "invalid Python source character",
147 ));
148 }
149 while self.indents.len() > 1 {
150 self.indents.pop();
151 self.emit_zero(TokenKind::Dedent)?;
152 }
153 self.emit_zero(TokenKind::End)?;
154 Ok(self.out)
155 }
156
157 fn layout(&mut self) -> Result<(), Diagnostic> {
158 let start = self.pos;
159 while matches!(self.peek(), Some(' ' | '\t' | '\u{0c}')) {
160 self.bump();
161 }
162 let prefix = &self.source[start..self.pos];
163 if self.peek().is_none() || matches!(self.peek(), Some('\r' | '\n' | '#')) {
164 if self.pos > start {
165 self.emit(TokenKind::Trivia, start)?;
166 }
167 self.at_line_start = false;
168 return Ok(());
169 }
170 let width = indent_width(prefix);
171 let (current, current_prefix) = self.indents.last().expect("base indent").clone();
172 if width == current
173 && prefix != current_prefix
174 && prefix.contains('\t') != current_prefix.contains('\t')
175 {
176 return Err(self.error(
177 Code::AmbiguousIndentation,
178 start,
179 "inconsistent tabs and spaces in indentation",
180 ));
181 }
182 if self.pos > start {
183 self.emit(TokenKind::Trivia, start)?;
184 }
185 if width > current {
186 if self.indents.len() >= self.limits.max_nesting {
187 return Err(self.error(
188 Code::ResourceLimit,
189 start,
190 "indentation nesting limit exceeded",
191 ));
192 }
193 self.indents.push((width, prefix.to_owned()));
194 self.emit_zero(TokenKind::Indent)?;
195 } else if width < current {
196 while self.indents.last().is_some_and(|(w, _)| *w > width) {
197 self.indents.pop();
198 self.emit_zero(TokenKind::Dedent)?;
199 }
200 if self.indents.last().map(|x| x.0) != Some(width) {
201 return Err(self.error(
202 Code::InvalidIndentation,
203 start,
204 "unindent does not match an outer level",
205 ));
206 }
207 }
208 self.at_line_start = false;
209 Ok(())
210 }
211
212 fn name_or_string(&mut self, start: usize) -> Result<(), Diagnostic> {
213 self.take_while(is_name_continue);
214 let word = &self.source[start..self.pos];
215 if self.peek().is_some_and(|c| matches!(c, '\'' | '"')) && is_string_prefix(word) {
216 let kind = if word.to_ascii_lowercase().contains('f') {
217 TokenKind::FString
218 } else if word.to_ascii_lowercase().contains('t') {
219 TokenKind::TemplateString
220 } else {
221 TokenKind::String
222 };
223 self.string(start, word, kind)
224 } else {
225 self.emit(
226 if KEYWORDS.contains(&word) {
227 TokenKind::Keyword
228 } else {
229 TokenKind::Name
230 },
231 start,
232 )
233 }
234 }
235
236 fn string(&mut self, start: usize, _prefix: &str, kind: TokenKind) -> Result<(), Diagnostic> {
237 let quote = self.peek().expect("string quote");
238 self.bump();
239 let triple = self.rest().starts_with(quote)
240 && self
241 .rest()
242 .get(quote.len_utf8()..)
243 .is_some_and(|s| s.starts_with(quote));
244 if triple {
245 self.bump();
246 self.bump();
247 }
248 let mut braces = 0usize;
249 loop {
250 let Some(ch) = self.peek() else {
251 return Err(self.error(
252 Code::UnterminatedLiteral,
253 start,
254 "unterminated string literal",
255 ));
256 };
257 if ch == '\\' {
258 self.bump();
259 if self.peek().is_some() {
260 self.bump();
261 }
262 continue;
263 }
264 if matches!(kind, TokenKind::FString | TokenKind::TemplateString) {
265 if ch == '{' && !self.rest().starts_with("{{") {
266 braces += 1;
267 if braces > self.limits.max_nesting {
268 return Err(self.error(
269 Code::ResourceLimit,
270 self.pos,
271 "interpolation nesting limit exceeded",
272 ));
273 }
274 }
275 if ch == '}' && !self.rest().starts_with("}}") {
276 if braces == 0 {
277 return Err(self.error(
278 Code::InvalidSyntax,
279 self.pos,
280 "single closing brace in interpolated string",
281 ));
282 }
283 braces -= 1;
284 }
285 }
286 if ch == quote {
287 self.bump();
288 if !triple
289 || (self.peek() == Some(quote)
290 && self
291 .rest()
292 .get(quote.len_utf8()..)
293 .is_some_and(|s| s.starts_with(quote)))
294 {
295 if triple {
296 self.bump();
297 self.bump();
298 }
299 if braces != 0 {
300 return Err(self.error(
301 Code::UnterminatedLiteral,
302 start,
303 "unterminated interpolation field",
304 ));
305 }
306 return self.emit(kind, start);
307 }
308 continue;
309 }
310 if !triple && matches!(ch, '\r' | '\n') {
311 return Err(self.error(
312 Code::UnterminatedLiteral,
313 start,
314 "unterminated string literal",
315 ));
316 }
317 self.bump();
318 }
319 }
320
321 fn number(&mut self, start: usize) -> Result<(), Diagnostic> {
322 self.take_while(|c| c.is_ascii_alphanumeric() || matches!(c, '_' | '.'));
323 if matches!(self.peek(), Some('+' | '-'))
324 && self.source[start..self.pos].ends_with(['e', 'E'])
325 {
326 self.bump();
327 self.take_while(|c| c.is_ascii_digit() || c == '_');
328 }
329 self.emit(TokenKind::Number, start)
330 }
331
332 fn newline(&mut self) -> Result<(), Diagnostic> {
333 let start = self.pos;
334 if self.peek() == Some('\r') {
335 self.bump();
336 }
337 if self.peek() == Some('\n') {
338 self.bump();
339 }
340 self.emit(TokenKind::Newline, start)?;
341 self.line += 1;
342 self.column = 0;
343 self.at_line_start = true;
344 Ok(())
345 }
346
347 fn emit(&mut self, kind: TokenKind, start: usize) -> Result<(), Diagnostic> {
348 let (line, column) = locate(self.source, start);
349 self.push(Token {
350 kind,
351 span: Span {
352 start,
353 end: self.pos,
354 },
355 line,
356 column,
357 })
358 }
359 fn emit_zero(&mut self, kind: TokenKind) -> Result<(), Diagnostic> {
360 self.push(Token {
361 kind,
362 span: Span {
363 start: self.pos,
364 end: self.pos,
365 },
366 line: self.line,
367 column: self.column,
368 })
369 }
370 fn push(&mut self, token: Token) -> Result<(), Diagnostic> {
371 if self.out.len() >= self.limits.max_tokens {
372 return Err(self.error(
373 Code::ResourceLimit,
374 token.span.start,
375 "token limit exceeded",
376 ));
377 }
378 self.out.push(token);
379 Ok(())
380 }
381 fn rest(&self) -> &str {
382 &self.source[self.pos..]
383 }
384 fn peek(&self) -> Option<char> {
385 self.rest().chars().next()
386 }
387 fn bump(&mut self) {
388 if let Some(c) = self.peek() {
389 self.pos += c.len_utf8();
390 self.column += 1;
391 }
392 }
393 fn take_while(&mut self, test: impl Fn(char) -> bool) {
394 while self.peek().is_some_and(&test) {
395 self.bump();
396 }
397 }
398 fn error(&self, code: Code, at: usize, message: &str) -> Diagnostic {
399 let (line, column) = locate(self.source, at);
400 Diagnostic {
401 code,
402 span: Span { start: at, end: at },
403 line,
404 column,
405 message: message.to_owned(),
406 }
407 }
408}
409
410fn is_name_start(c: char) -> bool {
411 c == '_' || c.is_alphabetic()
412}
413fn is_name_continue(c: char) -> bool {
414 c == '_' || c.is_alphanumeric()
415}
416fn is_string_prefix(s: &str) -> bool {
417 matches!(
418 s.to_ascii_lowercase().as_str(),
419 "r" | "u" | "b" | "br" | "rb" | "f" | "fr" | "rf" | "t" | "tr" | "rt"
420 )
421}
422fn indent_width(s: &str) -> usize {
423 s.chars().fold(0, |n, c| match c {
424 ' ' => n + 1,
425 '\t' => (n / 8 + 1) * 8,
426 '\u{0c}' => 0,
427 _ => n,
428 })
429}
430fn locate(source: &str, at: usize) -> (usize, usize) {
431 let before = &source[..at];
432 let line = before.bytes().filter(|b| *b == b'\n').count() + 1;
433 let col = before
434 .rsplit_once('\n')
435 .map_or(before, |(_, tail)| tail)
436 .chars()
437 .count();
438 (line, col)
439}