diff --git a/Cargo.lock b/Cargo.lock index ffece1cb3..1a291035b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3283,6 +3283,10 @@ dependencies = [ "pkcs8", "pymath", "rand_core 0.9.5", + "ruff_python_ast", + "ruff_python_parser", + "ruff_source_file", + "ruff_text_size", "rustix", "rustls", "rustls-native-certs", diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index a40a5bf24..7d230fcc0 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -31,6 +31,11 @@ rustpython-derive = { workspace = true } rustpython-vm = { workspace = true, default-features = false, features = ["compiler"]} rustpython-common = { workspace = true } +ruff_python_parser = { workspace = true } +ruff_python_ast = { workspace = true } +ruff_text_size = { workspace = true } +ruff_source_file = { workspace = true } + ahash = { workspace = true } ascii = { workspace = true } cfg-if = { workspace = true } diff --git a/crates/stdlib/src/lib.rs b/crates/stdlib/src/lib.rs index 8c234c22f..04aa623d1 100644 --- a/crates/stdlib/src/lib.rs +++ b/crates/stdlib/src/lib.rs @@ -49,6 +49,7 @@ mod pystruct; mod random; mod statistics; mod suggestions; +mod tokenize; // TODO: maybe make this an extension module, if we ever get those // mod re; #[cfg(all(feature = "host_env", not(target_arch = "wasm32")))] @@ -225,6 +226,7 @@ pub fn stdlib_module_defs(ctx: &Context) -> Vec<&'static builtins::PyModuleDef> ssl::module_def(ctx), statistics::module_def(ctx), suggestions::module_def(ctx), + tokenize::module_def(ctx), #[cfg(all(feature = "host_env", unix, not(target_os = "redox")))] syslog::module_def(ctx), #[cfg(all( diff --git a/crates/stdlib/src/tokenize.rs b/crates/stdlib/src/tokenize.rs new file mode 100644 index 000000000..33667a203 --- /dev/null +++ b/crates/stdlib/src/tokenize.rs @@ -0,0 +1,391 @@ +pub(crate) use _tokenize::module_def; + +#[pymodule] +mod _tokenize { + use crate::{ + common::lock::PyRwLock, + vm::{ + AsObject, Py, PyPayload, PyResult, VirtualMachine, + builtins::{PyBytes, PyStr, PyType}, + convert::ToPyObject, + function::ArgCallable, + protocol::PyIterReturn, + types::{Constructor, IterNext, Iterable, SelfIter}, + }, + }; + use ruff_python_ast::PySourceType; + use ruff_python_ast::token::{Token, TokenKind, Tokens}; + use ruff_python_parser::{ParseError, parse_unchecked_source}; + use ruff_source_file::{LineIndex, LineRanges}; + use ruff_text_size::{Ranged, TextRange}; + use std::{cmp::Ordering, fmt}; + + /// `__import__("token").OP` + const TOKEN_OP: u8 = 55; + + #[pyattr] + #[pyclass(name = "TokenizerIter")] + #[derive(PyPayload)] + pub struct PyTokenizerIter { + readline: ArgCallable, // TODO: This should be PyObject + extra_tokens: bool, + encoding: Option, + state: PyRwLock, + } + + impl PyTokenizerIter { + fn readline(&self, vm: &VirtualMachine) -> PyResult { + // TODO: When `readline` is PyObject, + // we need to check if it's callable and raise a type error if it's not. + let raw_line = match self.readline.invoke((), vm) { + Ok(v) => v, + Err(err) => { + if err.fast_isinstance(vm.ctx.exceptions.stop_iteration) { + return Ok(String::new()); + } + return Err(err); + } + }; + Ok(match &self.encoding { + Some(encoding) => { + let bytes = raw_line + .downcast::() + .map_err(|_| vm.new_type_error("readline() returned a non-bytes object"))?; + vm.state + .codec_registry + .decode_text(bytes.into(), encoding, None, vm) + .map(|s| s.to_string())? + } + None => raw_line + .downcast::() + .map(|s| s.to_string()) + .map_err(|_| vm.new_type_error("readline() returned a non-string object"))?, + }) + } + } + + impl fmt::Debug for PyTokenizerIter { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PyTokenizerIter") + .field("readline", &self.readline) + .field("encoding", &self.encoding) + .field("extra_tokens", &self.extra_tokens) + .finish() + } + } + + #[pyclass(with(Constructor, Iterable, IterNext))] + impl PyTokenizerIter {} + + impl Constructor for PyTokenizerIter { + type Args = PyTokenizerIterArgs; + + fn py_new(_cls: &Py, args: Self::Args, _vm: &VirtualMachine) -> PyResult { + let Self::Args { + readline, + extra_tokens, + encoding, + } = args; + + Ok(Self { + readline, + extra_tokens, + encoding: encoding.map(|s| s.to_string()), + state: PyRwLock::new(PyTokenizerIterState::default()), + }) + } + } + + impl SelfIter for PyTokenizerIter {} + + impl IterNext for PyTokenizerIter { + fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { + let mut state = { + let guard = zelf.state.read(); + guard.clone() + }; + + if state.eof { + return Ok(PyIterReturn::StopIteration(None)); + } + + let token = loop { + // TODO: Check here for errors. Raise SyntaxError if needed + + if let Some(tok) = state.next_token() { + break tok; + } + + let nline = zelf.readline(vm)?; + if nline.is_empty() { + state.eof = true; + *zelf.state.write() = state.clone(); + + let line_num = &state.start().0; + let out = vm + .ctx + .new_tuple(vec![ + token_kind_value(TokenKind::EndOfFile).to_pyobject(vm), + vm.ctx.new_str("").into(), + vm.ctx + .new_tuple(vec![line_num.to_pyobject(vm), (-1).to_pyobject(vm)]) + .into(), + vm.ctx + .new_tuple(vec![line_num.to_pyobject(vm), (-1).to_pyobject(vm)]) + .into(), + vm.ctx.new_str(state.current_line()).into(), + ]) + .into(); + return Ok(PyIterReturn::Return(out)); + } + state.push_line(&nline); + }; + + *zelf.state.write() = state.clone(); + + let token_kind = token.kind(); + let token_value = if zelf.extra_tokens && token_kind.is_operator() { + TOKEN_OP + } else { + token_kind_value(token_kind) + }; + let (start_x, start_y) = &state.start(); + let (end_x, end_y) = &state.end(); + + let mut token_repr = &state.source[state.range()]; + if !zelf.extra_tokens { + token_repr = token_repr.trim(); + } + + let out = vm + .ctx + .new_tuple(vec![ + token_value.to_pyobject(vm), + vm.ctx.new_str(token_repr).into(), + vm.ctx + .new_tuple(vec![start_x.to_pyobject(vm), start_y.to_pyobject(vm)]) + .into(), + vm.ctx + .new_tuple(vec![end_x.to_pyobject(vm), end_y.to_pyobject(vm)]) + .into(), + vm.ctx.new_str(state.current_line()).into(), + ]) + .into(); + Ok(PyIterReturn::Return(out)) + } + } + + #[derive(FromArgs)] + pub struct PyTokenizerIterArgs { + #[pyarg(positional)] + readline: ArgCallable, + #[pyarg(named)] + extra_tokens: bool, + #[pyarg(named, optional)] + encoding: Option>, + } + + #[derive(Clone, Debug)] + struct PyTokenizerIterState { + /// Source code. + source: String, + prev_token: Option, + /// Tokens of `source`. + tokens: Tokens, + /// Errors of `source` + errors: Vec, + /// LineIndex of `source`. + line_index: LineIndex, + /// Marker that says we already emitted EOF, and needs to stop iterating. + eof: bool, + } + + impl PyTokenizerIterState { + fn push_line(&mut self, line: &str) { + self.source.push_str(line); + + let parsed = parse_unchecked_source(&self.source, PySourceType::Python); + self.tokens = parsed.tokens().clone(); + self.errors = parsed.errors().to_vec(); + self.line_index = LineIndex::from_source_text(&self.source); + } + + #[must_use] + fn current_line(&self) -> &str { + let (kind, range): (TokenKind, TextRange) = match self.prev_token { + Some(token) => token.as_tuple(), + None => (TokenKind::Unknown, TextRange::default()), + }; + + match kind { + TokenKind::Newline => self.source.full_line_str(range.start()), + _ => self.source.full_lines_str(range), + } + } + + #[must_use] + fn next_token(&mut self) -> Option { + for token in self.tokens.iter() { + let (kind, range): (TokenKind, TextRange) = token.as_tuple(); + + if matches!(kind, TokenKind::NonLogicalNewline) { + continue; + } + + if matches!(range.ordering(self.range()), Ordering::Greater) { + self.prev_token = Some(*token); + return self.prev_token; + } + } + + None + } + + #[must_use] + fn range(&self) -> TextRange { + match self.prev_token { + Some(token) => token.range(), + None => TextRange::default(), + } + } + + #[must_use] + fn start(&self) -> (usize, usize) { + let lc = self + .line_index + .line_column(self.range().start(), &self.source); + (lc.line.get(), lc.column.to_zero_indexed()) + } + + #[must_use] + fn end(&self) -> (usize, usize) { + let lc = self + .line_index + .line_column(self.range().end(), &self.source); + (lc.line.get(), lc.column.to_zero_indexed()) + } + } + + impl Default for PyTokenizerIterState { + fn default() -> Self { + const SOURCE: &str = ""; + let parsed = parse_unchecked_source(SOURCE, PySourceType::Python); + + Self { + source: SOURCE.to_owned(), + prev_token: None, + tokens: parsed.tokens().clone(), + errors: parsed.errors().to_vec(), + line_index: LineIndex::from_source_text(SOURCE), + eof: false, + } + } + } + + const fn token_kind_value(kind: TokenKind) -> u8 { + match kind { + TokenKind::EndOfFile => 0, + TokenKind::Name + | TokenKind::For + | TokenKind::In + | TokenKind::Pass + | TokenKind::Class + | TokenKind::And + | TokenKind::Is + | TokenKind::Raise + | TokenKind::True + | TokenKind::False + | TokenKind::Assert + | TokenKind::Try + | TokenKind::While + | TokenKind::Yield + | TokenKind::Lambda + | TokenKind::None + | TokenKind::Not + | TokenKind::Or + | TokenKind::Break + | TokenKind::Continue + | TokenKind::Global + | TokenKind::Nonlocal + | TokenKind::Return + | TokenKind::Except + | TokenKind::Import + | TokenKind::Case + | TokenKind::Match + | TokenKind::Type + | TokenKind::Await + | TokenKind::With + | TokenKind::Del + | TokenKind::Finally + | TokenKind::From + | TokenKind::Def + | TokenKind::If + | TokenKind::Else + | TokenKind::Elif + | TokenKind::As + | TokenKind::Async => 1, + TokenKind::Int | TokenKind::Complex | TokenKind::Float => 2, + TokenKind::String => 3, + TokenKind::Newline | TokenKind::NonLogicalNewline => 4, + TokenKind::Indent => 5, + TokenKind::Dedent => 6, + TokenKind::Lpar => 7, + TokenKind::Rpar => 8, + TokenKind::Lsqb => 9, + TokenKind::Rsqb => 10, + TokenKind::Colon => 11, + TokenKind::Comma => 12, + TokenKind::Semi => 13, + TokenKind::Plus => 14, + TokenKind::Minus => 15, + TokenKind::Star => 16, + TokenKind::Slash => 17, + TokenKind::Vbar => 18, + TokenKind::Amper => 19, + TokenKind::Less => 20, + TokenKind::Greater => 21, + TokenKind::Equal => 22, + TokenKind::Dot => 23, + TokenKind::Percent => 24, + TokenKind::Lbrace => 25, + TokenKind::Rbrace => 26, + TokenKind::EqEqual => 27, + TokenKind::NotEqual => 28, + TokenKind::LessEqual => 29, + TokenKind::GreaterEqual => 30, + TokenKind::Tilde => 31, + TokenKind::CircumFlex => 32, + TokenKind::LeftShift => 33, + TokenKind::RightShift => 34, + TokenKind::DoubleStar => 35, + TokenKind::PlusEqual => 36, + TokenKind::MinusEqual => 37, + TokenKind::StarEqual => 38, + TokenKind::SlashEqual => 39, + TokenKind::PercentEqual => 40, + TokenKind::AmperEqual => 41, + TokenKind::VbarEqual => 42, + TokenKind::CircumflexEqual => 43, + TokenKind::LeftShiftEqual => 44, + TokenKind::RightShiftEqual => 45, + TokenKind::DoubleStarEqual => 46, + TokenKind::DoubleSlash => 47, + TokenKind::DoubleSlashEqual => 48, + TokenKind::At => 49, + TokenKind::AtEqual => 50, + TokenKind::Rarrow => 51, + TokenKind::Ellipsis => 52, + TokenKind::ColonEqual => 53, + TokenKind::Exclamation => 54, + TokenKind::FStringStart => 59, + TokenKind::FStringMiddle => 60, + TokenKind::FStringEnd => 61, + TokenKind::Comment => 62, + TokenKind::TStringStart => 62, // 3.14 compatible + TokenKind::TStringMiddle => 63, // 3.14 compatible + TokenKind::TStringEnd => 64, // 3.14 compatible + TokenKind::IpyEscapeCommand | TokenKind::Question => 0, // Ruff's specific + TokenKind::Unknown => 0, + } + } +}