1use std::fmt::Formatter;
2use std::sync::Arc;
3
4use arc_swap::ArcSwapOption;
5use get_size2::GetSize;
6use ruff_python_ast::{
7 AnyRootNodeRef, HasNodeIndex, ModExpression, ModModule, NodeIndex, NodeIndexError,
8 PythonVersion, StringLiteral,
9};
10use ruff_python_parser::{
11 ParseError, ParseErrorType, ParseOptions, Parsed, parse_cells_unchecked,
12 parse_string_annotation, parse_unchecked,
13};
14
15use crate::files::File;
16use crate::source::source_text;
17use crate::{Db, PythonFile};
18
19#[salsa::tracked(returns(ref), no_eq, heap_size=ruff_memory_usage::heap_size, lru=200)]
35pub fn parsed_module(db: &dyn Db, file: PythonFile<'_>) -> ParsedModule {
36 let source_file = file.file(db);
37 let python_version = file.python_version(db);
38 let _span = tracing::trace_span!("parsed_module", ?source_file, %python_version).entered();
39
40 let parsed = parsed_module_impl(db, source_file, python_version);
41
42 ParsedModule::new(source_file, python_version, parsed)
43}
44
45pub(super) fn disable_lru(db: &mut dyn Db) {
46 parsed_module::set_lru_capacity(db, 0);
47}
48
49fn parsed_module_impl(db: &dyn Db, file: File, target_version: PythonVersion) -> Parsed<ModModule> {
50 let source = source_text(db, file);
51 let ty = file.source_type(db);
52
53 let options = ParseOptions::from(ty).with_target_version(target_version);
54
55 if let Some(notebook) = source.as_notebook() {
59 parse_cells_unchecked(&source, notebook.cell_offsets().content_ranges(), &options)
60 } else {
61 parse_unchecked(&source, options)
62 .try_into_module()
63 .expect("PySourceType always parses into a module")
64 }
65}
66
67pub fn parsed_string_annotation(
68 source: &str,
69 string: &StringLiteral,
70) -> Result<Parsed<ModExpression>, ParseError> {
71 let expr = parse_string_annotation(source, string)?;
72
73 indexed::ensure_indexed(&expr, string.node_index().load()).map_err(|err| {
75 let message = match err {
76 NodeIndexError::NoParent => {
77 "Internal error: string annotation's parent had no NodeIndex"
78 }
79 NodeIndexError::TooNested => {
80 "Too many levels of nested string annotations; \
81 remove the redundant nested quotes"
82 }
83 NodeIndexError::OverflowedIndices => {
84 "File too long for string annotations; either break up the file \
85 or don't use string annotations"
86 }
87 NodeIndexError::OverflowedSubIndices => {
88 "File too long for nested string annotations; remove the redundant nested quotes"
89 }
90 NodeIndexError::ExhaustedSubIndices => {
91 "String annotation is too long; consider introducing type aliases to simplify"
92 }
93 NodeIndexError::ExhaustedSubSubIndices => {
94 "Nested string annotation is too long; remove the redundant nested quotes"
95 }
96 };
97
98 ParseError {
99 error: ParseErrorType::StringAnnotationError(message),
100 location: string.range,
101 }
102 })?;
103
104 Ok(expr)
105}
106
107#[derive(Clone, get_size2::GetSize)]
112pub struct ParsedModule {
113 file: File,
114 python_version: PythonVersion,
115 #[get_size(size_fn = arc_swap_size)]
116 inner: Arc<ArcSwapOption<indexed::IndexedModule>>,
117}
118
119impl ParsedModule {
120 pub fn new(file: File, python_version: PythonVersion, parsed: Parsed<ModModule>) -> Self {
121 Self {
122 file,
123 python_version,
124 inner: Arc::new(ArcSwapOption::new(Some(indexed::IndexedModule::new(
125 parsed,
126 )))),
127 }
128 }
129 pub fn load(&self, db: &dyn Db) -> ParsedModuleRef {
134 let parsed = match self.inner.load_full() {
135 Some(parsed) => parsed,
136 None => {
137 let parsed = indexed::IndexedModule::new(parsed_module_impl(
139 db,
140 self.file,
141 self.python_version,
142 ));
143 tracing::debug!(
144 "File `{}` was reparsed after being collected in the current Salsa revision",
145 self.file.path(db)
146 );
147
148 self.inner.store(Some(parsed.clone()));
149 parsed
150 }
151 };
152
153 ParsedModuleRef {
154 module: self.clone(),
155 indexed: parsed,
156 }
157 }
158
159 pub fn clear(&self) {
161 self.inner.store(None);
162 }
163
164 pub fn file(&self) -> File {
166 self.file
167 }
168
169 pub fn python_version(&self) -> PythonVersion {
171 self.python_version
172 }
173}
174
175impl std::fmt::Debug for ParsedModule {
176 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
177 f.debug_tuple("ParsedModule").field(&self.inner).finish()
178 }
179}
180
181impl PartialEq for ParsedModule {
182 fn eq(&self, other: &Self) -> bool {
183 Arc::ptr_eq(&self.inner, &other.inner)
184 }
185}
186
187impl Eq for ParsedModule {}
188
189#[derive(Clone)]
191pub struct ParsedModuleRef {
192 module: ParsedModule,
193 indexed: Arc<indexed::IndexedModule>,
194}
195
196impl ParsedModuleRef {
197 pub fn module(&self) -> &ParsedModule {
199 &self.module
200 }
201
202 pub fn get_by_index<'ast>(&'ast self, index: NodeIndex) -> AnyRootNodeRef<'ast> {
204 self.indexed.get_by_index(index)
205 }
206}
207
208impl std::ops::Deref for ParsedModuleRef {
209 type Target = Parsed<ModModule>;
210
211 fn deref(&self) -> &Self::Target {
212 &self.indexed.parsed
213 }
214}
215
216fn arc_swap_size<T>(arc_swap: &Arc<ArcSwapOption<T>>) -> usize
218where
219 T: GetSize,
220{
221 if let Some(value) = &*arc_swap.load() {
222 T::get_heap_size(value)
223 } else {
224 0
225 }
226}
227
228mod indexed {
229 use std::sync::Arc;
230
231 use ruff_python_ast::visitor::source_order::*;
232 use ruff_python_ast::*;
233 use ruff_python_parser::Parsed;
234
235 #[derive(Debug, get_size2::GetSize)]
237 pub struct IndexedModule {
238 index: IndexedNodes,
239 pub parsed: Parsed<ModModule>,
240 }
241
242 #[derive(Debug, Default)]
269 struct IndexedNodes {
270 chunks: Box<[IndexChunk]>,
271 words: Box<[u64]>,
272 }
273
274 #[derive(Debug, get_size2::GetSize)]
276 struct IndexChunk {
277 base: usize,
279 word_start: u32,
281 entry_bits: u8,
283 entry_count: u8,
286 layout: IndexChunkLayout,
287 }
288
289 impl get_size2::GetSize for IndexedNodes {
290 fn get_heap_size_with_tracker<T: get_size2::GetSizeTracker>(
291 &self,
292 tracker: T,
293 ) -> (usize, T) {
294 let (chunks_size, tracker) =
295 get_size2::GetSize::get_heap_size_with_tracker(&self.chunks, tracker);
296 let words = self
297 .chunks
298 .iter()
299 .map(|chunk| {
300 (usize::from(chunk.entry_count) * Self::REPORTED_ENTRY_BITS)
301 .div_ceil(u64::BITS as usize)
302 })
303 .sum::<usize>();
304
305 (chunks_size + words * size_of::<u64>(), tracker)
306 }
307 }
308
309 #[derive(Copy, Clone, Debug, get_size2::GetSize)]
310 #[repr(u8)]
311 enum IndexChunkLayout {
312 Relative,
322 Wide,
329 }
330
331 #[derive(Default)]
332 struct IndexedNodesBuilder<'ast> {
333 chunks: Vec<IndexChunk>,
334 words: Vec<u64>,
335 pending: Vec<AnyRootNodeRef<'ast>>,
336 #[cfg(test)]
337 all_nodes: Vec<AnyRootNodeRef<'ast>>,
338 }
339
340 impl<'ast> IndexedNodesBuilder<'ast> {
341 fn new() -> Self {
342 Self {
343 pending: Vec::with_capacity(IndexedNodes::CHUNK_LEN),
344 ..Self::default()
345 }
346 }
347
348 fn push(&mut self, node: AnyRootNodeRef<'ast>) {
349 #[cfg(test)]
350 self.all_nodes.push(node);
351
352 self.pending.push(node);
353
354 if self.pending.len() == IndexedNodes::CHUNK_LEN {
355 self.flush();
356 }
357 }
358
359 fn finish(mut self) -> IndexedNodes {
360 self.flush();
361
362 IndexedNodes {
363 chunks: self.chunks.into_boxed_slice(),
364 words: self.words.into_boxed_slice(),
365 }
366 }
367
368 fn flush(&mut self) {
369 IndexedNodes::extend_from_nodes(&mut self.chunks, &mut self.words, &self.pending);
370 self.pending.clear();
371 }
372 }
373
374 impl IndexedNodes {
375 const ALIGNMENT: usize = std::mem::align_of::<AtomicNodeIndex>();
376 const CHUNK_LEN: usize = 64;
377 const KIND_BITS: u8 = 5;
378 const KIND_MASK: u64 = (1 << Self::KIND_BITS) - 1;
379 const REPORTED_ENTRY_BITS: usize = 32;
380
381 fn extend_from_nodes(
382 chunks: &mut Vec<IndexChunk>,
383 words: &mut Vec<u64>,
384 nodes: &[AnyRootNodeRef<'_>],
385 ) {
386 for node_chunk in nodes.chunks(Self::CHUNK_LEN) {
387 let (base, max, aligned) =
388 node_chunk
389 .iter()
390 .fold((usize::MAX, 0, true), |(base, max, aligned), node| {
391 let (_, pointer) = node.into_raw_parts();
392 let address = pointer.as_ptr().expose_provenance();
393 (
394 base.min(address),
395 max.max(address),
396 aligned && address.is_multiple_of(Self::ALIGNMENT),
397 )
398 });
399 let offset_bits = usize::BITS - ((max - base) / Self::ALIGNMENT).leading_zeros();
400 let relative_bits = u8::try_from(offset_bits)
401 .expect("an address offset cannot require more than u8::MAX bits")
402 + Self::KIND_BITS;
403 let word_start = u32::try_from(words.len())
404 .expect("indexed AST bitstream should fit in u32 words");
405
406 if aligned && relative_bits <= 64 {
407 let entry_count = u8::try_from(node_chunk.len())
408 .expect("an index chunk contains at most 64 entries");
409 chunks.push(IndexChunk {
410 base,
411 word_start,
412 entry_bits: relative_bits,
413 entry_count,
414 layout: IndexChunkLayout::Relative,
415 });
416 for (entry, node) in node_chunk.iter().enumerate() {
417 let (kind, pointer) = node.into_raw_parts();
418 let address = pointer.as_ptr().expose_provenance();
419 let offset = (address - base) / Self::ALIGNMENT;
420 let offset = u64::try_from(offset)
421 .expect("relative address offset was checked to fit in 64 bits");
422 Self::write_bits(
423 words,
424 word_start as usize * 64 + entry * usize::from(relative_bits),
425 (offset << Self::KIND_BITS) | u64::from(kind as u8),
426 relative_bits,
427 );
428 }
429 } else {
430 let entry_count = u8::try_from(node_chunk.len())
432 .expect("an index chunk contains at most 64 entries");
433 chunks.push(IndexChunk {
434 base: 0,
435 word_start,
436 entry_bits: Self::KIND_BITS,
437 entry_count,
438 layout: IndexChunkLayout::Wide,
439 });
440 words.extend(node_chunk.iter().map(|node| {
441 let (_, pointer) = node.into_raw_parts();
442 u64::try_from(pointer.as_ptr().expose_provenance())
443 .expect("AST node addresses should fit in a bitstream word")
444 }));
445 for (entry, node) in node_chunk.iter().enumerate() {
446 let (kind, _) = node.into_raw_parts();
447 Self::write_bits(
448 words,
449 (word_start as usize + node_chunk.len()) * 64
450 + entry * usize::from(Self::KIND_BITS),
451 u64::from(kind as u8),
452 Self::KIND_BITS,
453 );
454 }
455 }
456 }
457 }
458
459 fn write_bits(words: &mut Vec<u64>, bit: usize, value: u64, bits: u8) {
460 debug_assert!((1..=64).contains(&bits));
461 let word = bit / 64;
462 let shift = bit % 64;
463 let end = bit + usize::from(bits);
464 words.resize(words.len().max(end.div_ceil(64)), 0);
465 words[word] |= value << shift;
466 if end > (word + 1) * 64 {
467 words[word + 1] |= value >> (64 - shift);
468 }
469 }
470
471 fn read_bits(words: &[u64], bit: usize, bits: u8) -> u64 {
472 debug_assert!((1..=64).contains(&bits));
473 let word = bit / 64;
474 let shift = bit % 64;
475 let low = words[word] >> shift;
476 let value = if shift + usize::from(bits) <= 64 {
477 low
478 } else {
479 low | (words[word + 1] << (64 - shift))
480 };
481 if bits == 64 {
482 value
483 } else {
484 value & ((1 << bits) - 1)
485 }
486 }
487
488 #[cfg(test)]
489 fn len(&self) -> usize {
490 self.chunks
491 .iter()
492 .map(|chunk| usize::from(chunk.entry_count))
493 .sum()
494 }
495
496 fn get(&self, index: usize) -> (usize, RootNodeKind) {
497 let chunk_index = index / Self::CHUNK_LEN;
498 let entry_index = index % Self::CHUNK_LEN;
499 let chunk = &self.chunks[chunk_index];
500
501 assert!(
503 entry_index < usize::from(chunk.entry_count),
504 "index out of bounds: the len is {} but the index is {index}",
505 chunk_index * Self::CHUNK_LEN + usize::from(chunk.entry_count),
506 );
507
508 let words = &self.words[chunk.word_start as usize..];
509
510 match chunk.layout {
511 IndexChunkLayout::Relative => {
512 let entry = Self::read_bits(
513 words,
514 entry_index * usize::from(chunk.entry_bits),
515 chunk.entry_bits,
516 );
517 let offset = (entry >> Self::KIND_BITS) as usize;
518 let kind = RootNodeKind::from_u8((entry & Self::KIND_MASK) as u8)
519 .expect("packed node kind should be valid");
520 (chunk.base + offset * Self::ALIGNMENT, kind)
521 }
522 IndexChunkLayout::Wide => {
523 let address = usize::try_from(words[entry_index])
524 .expect("stored AST node address should fit in usize");
525 let kind_bit = usize::from(chunk.entry_count) * 64
526 + entry_index * usize::from(Self::KIND_BITS);
527 let kind =
528 RootNodeKind::from_u8(
529 Self::read_bits(words, kind_bit, Self::KIND_BITS) as u8
530 )
531 .expect("packed node kind should be valid");
532 (address, kind)
533 }
534 }
535 }
536 }
537
538 const _: () = assert!(RootNodeKind::ALL.len() <= 1 << IndexedNodes::KIND_BITS);
539
540 pub fn ensure_indexed(
543 parsed: &Parsed<ModExpression>,
544 parent_node_index: NodeIndex,
545 ) -> Result<(), NodeIndexError> {
546 let parent_index = parent_node_index.as_u32().ok_or(NodeIndexError::NoParent)?;
547 let (index, max_index) = sub_indices(parent_index)?;
548 let mut visitor = Visitor {
549 overflowed: false,
550 nodes: None,
551 index,
552 max_index,
553 };
554
555 AnyNodeRef::from(parsed.syntax()).visit_source_order(&mut visitor);
556
557 if visitor.overflowed {
558 let level = sub_ast_level(parent_index);
559 if level == 0 {
560 return Err(NodeIndexError::ExhaustedSubIndices);
561 } else {
562 return Err(NodeIndexError::ExhaustedSubSubIndices);
563 }
564 }
565
566 Ok(())
567 }
568
569 impl IndexedModule {
570 pub fn new(parsed: Parsed<ModModule>) -> Arc<Self> {
572 let mut visitor = Visitor {
573 nodes: Some(IndexedNodesBuilder::new()),
574 index: 0,
575 max_index: MAX_REAL_INDEX,
576 overflowed: false,
577 };
578
579 let mut inner = Arc::new(IndexedModule {
580 parsed,
581 index: IndexedNodes::default(),
582 });
583
584 AnyNodeRef::from(inner.parsed.syntax()).visit_source_order(&mut visitor);
585
586 let index = visitor
587 .nodes
588 .expect("top-level AST visitor should collect indexed nodes")
589 .finish();
590 Arc::get_mut(&mut inner)
591 .expect("newly created indexed module should have a unique Arc")
592 .index = index;
593
594 inner
595 }
596
597 pub fn get_by_index<'ast>(&'ast self, index: NodeIndex) -> AnyRootNodeRef<'ast> {
599 let index = index
600 .as_u32()
601 .expect("attempted to access uninitialized `NodeIndex`");
602
603 let index = index as usize;
604 let (address, kind) = self.index.get(index);
605
606 unsafe {
610 AnyRootNodeRef::from_raw_parts(
611 kind,
612 std::ptr::NonNull::with_exposed_provenance(
613 std::num::NonZeroUsize::new(address)
614 .expect("recorded AST node address should be non-null"),
615 ),
616 )
617 }
618 }
619 }
620
621 struct Visitor<'ast> {
623 index: u32,
624 max_index: u32,
625 nodes: Option<IndexedNodesBuilder<'ast>>,
626 overflowed: bool,
627 }
628
629 impl<'ast> Visitor<'ast> {
630 fn visit_node<T>(&mut self, node: &'ast T)
631 where
632 T: HasNodeIndex,
633 AnyRootNodeRef<'ast>: From<&'ast T>,
634 {
635 if self.index > self.max_index {
637 self.overflowed = true;
638 } else {
639 node.node_index().set(NodeIndex::from(self.index));
640 }
641
642 if let Some(nodes) = &mut self.nodes {
643 nodes.push(AnyRootNodeRef::from(node));
644 }
645 self.index += 1;
646 }
647 }
648
649 impl<'a> SourceOrderVisitor<'a> for Visitor<'a> {
650 #[inline]
651 fn visit_stmt(&mut self, stmt: &'a Stmt) {
652 self.visit_node(stmt);
653 walk_stmt(self, stmt);
654 }
655
656 #[inline]
657 fn visit_annotation(&mut self, expr: &'a Expr) {
658 walk_annotation(self, expr);
660 }
661
662 #[inline]
663 fn visit_expr(&mut self, expr: &'a Expr) {
664 self.visit_node(expr);
665 walk_expr(self, expr);
666 }
667
668 #[inline]
669 fn visit_decorator(&mut self, decorator: &'a Decorator) {
670 self.visit_node(decorator);
671 walk_decorator(self, decorator);
672 }
673
674 #[inline]
675 fn visit_comprehension(&mut self, comprehension: &'a Comprehension) {
676 self.visit_node(comprehension);
677 walk_comprehension(self, comprehension);
678 }
679
680 #[inline]
681 fn visit_except_handler(&mut self, except_handler: &'a ExceptHandler) {
682 self.visit_node(except_handler);
683 walk_except_handler(self, except_handler);
684 }
685
686 #[inline]
687 fn visit_arguments(&mut self, arguments: &'a Arguments) {
688 self.visit_node(arguments);
689 walk_arguments(self, arguments);
690 }
691
692 #[inline]
693 fn visit_parameters(&mut self, parameters: &'a Parameters) {
694 self.visit_node(parameters);
695 walk_parameters(self, parameters);
696 }
697
698 #[inline]
699 fn visit_parameter(&mut self, arg: &'a Parameter) {
700 self.visit_node(arg);
701 walk_parameter(self, arg);
702 }
703
704 fn visit_parameter_with_default(
705 &mut self,
706 parameter_with_default: &'a ParameterWithDefault,
707 ) {
708 self.visit_node(parameter_with_default);
709 walk_parameter_with_default(self, parameter_with_default);
710 }
711
712 #[inline]
713 fn visit_keyword(&mut self, keyword: &'a Keyword) {
714 self.visit_node(keyword);
715 walk_keyword(self, keyword);
716 }
717
718 #[inline]
719 fn visit_alias(&mut self, alias: &'a Alias) {
720 self.visit_node(alias);
721 walk_alias(self, alias);
722 }
723
724 #[inline]
725 fn visit_with_item(&mut self, with_item: &'a WithItem) {
726 self.visit_node(with_item);
727 walk_with_item(self, with_item);
728 }
729
730 #[inline]
731 fn visit_type_params(&mut self, type_params: &'a TypeParams) {
732 self.visit_node(type_params);
733 walk_type_params(self, type_params);
734 }
735
736 #[inline]
737 fn visit_type_param(&mut self, type_param: &'a TypeParam) {
738 self.visit_node(type_param);
739 walk_type_param(self, type_param);
740 }
741
742 #[inline]
743 fn visit_match_case(&mut self, match_case: &'a MatchCase) {
744 self.visit_node(match_case);
745 walk_match_case(self, match_case);
746 }
747
748 #[inline]
749 fn visit_pattern(&mut self, pattern: &'a Pattern) {
750 self.visit_node(pattern);
751 walk_pattern(self, pattern);
752 }
753
754 #[inline]
755 fn visit_pattern_arguments(&mut self, pattern_arguments: &'a PatternArguments) {
756 self.visit_node(pattern_arguments);
757 walk_pattern_arguments(self, pattern_arguments);
758 }
759
760 #[inline]
761 fn visit_pattern_keyword(&mut self, pattern_keyword: &'a PatternKeyword) {
762 self.visit_node(pattern_keyword);
763 walk_pattern_keyword(self, pattern_keyword);
764 }
765
766 #[inline]
767 fn visit_elif_else_clause(&mut self, elif_else_clause: &'a ElifElseClause) {
768 self.visit_node(elif_else_clause);
769 walk_elif_else_clause(self, elif_else_clause);
770 }
771
772 #[inline]
773 fn visit_f_string(&mut self, f_string: &'a FString) {
774 self.visit_node(f_string);
775 walk_f_string(self, f_string);
776 }
777
778 #[inline]
779 fn visit_interpolated_string_element(
780 &mut self,
781 interpolated_string_element: &'a InterpolatedStringElement,
782 ) {
783 self.visit_node(interpolated_string_element);
784 walk_interpolated_string_element(self, interpolated_string_element);
785 }
786
787 #[inline]
788 fn visit_t_string(&mut self, t_string: &'a TString) {
789 self.visit_node(t_string);
790 walk_t_string(self, t_string);
791 }
792
793 #[inline]
794 fn visit_string_literal(&mut self, string_literal: &'a StringLiteral) {
795 self.visit_node(string_literal);
796 walk_string_literal(self, string_literal);
797 }
798
799 #[inline]
800 fn visit_bytes_literal(&mut self, bytes_literal: &'a BytesLiteral) {
801 self.visit_node(bytes_literal);
802 walk_bytes_literal(self, bytes_literal);
803 }
804
805 #[inline]
806 fn visit_identifier(&mut self, identifier: &'a Identifier) {
807 self.visit_node(identifier);
808 walk_identifier(self, identifier);
809 }
810 }
811
812 #[cfg(test)]
813 mod tests {
814 use super::*;
815
816 #[test]
817 #[should_panic(expected = "index out of bounds: the len is 1 but the index is 1")]
818 fn indexed_nodes_relative_tail_bounds() {
819 let index = IndexedNodes {
820 chunks: vec![IndexChunk {
821 base: 0x1000,
822 word_start: 0,
823 entry_bits: IndexedNodes::KIND_BITS,
824 entry_count: 1,
825 layout: IndexChunkLayout::Relative,
826 }]
827 .into_boxed_slice(),
828 words: vec![u64::from(RootNodeKind::Stmt as u8)].into_boxed_slice(),
829 };
830
831 assert_eq!(index.get(0), (0x1000, RootNodeKind::Stmt));
832 index.get(1);
833 }
834
835 #[test]
836 #[should_panic(expected = "index out of bounds: the len is 1 but the index is 1")]
837 fn indexed_nodes_wide_tail_bounds() {
838 let index = IndexedNodes {
839 chunks: vec![IndexChunk {
840 base: 0,
841 word_start: 0,
842 entry_bits: IndexedNodes::KIND_BITS,
843 entry_count: 1,
844 layout: IndexChunkLayout::Wide,
845 }]
846 .into_boxed_slice(),
847 words: vec![0x1000, u64::from(RootNodeKind::Stmt as u8)].into_boxed_slice(),
848 };
849
850 assert_eq!(index.get(0), (0x1000, RootNodeKind::Stmt));
851 index.get(1);
852 }
853
854 #[test]
855 #[should_panic(expected = "index out of bounds: the len is 65 but the index is 65")]
856 fn indexed_nodes_chunk_boundary() {
857 let parsed = ruff_python_parser::parse_module(&"pass\n".repeat(65)).unwrap();
858 let indexed = IndexedModule::new(parsed);
859
860 assert_eq!(indexed.index.len(), 65);
861 assert_eq!(indexed.index.get(63).1, RootNodeKind::Stmt);
862 assert_eq!(indexed.index.get(64).1, RootNodeKind::Stmt);
863 indexed.index.get(65);
864 }
865
866 #[test]
867 fn indexed_nodes_round_trip() {
868 let parsed = ruff_python_parser::parse_module(
869 r#"
870import os as imported_os
871
872@decorator
873class C[T](Base, metaclass=Meta):
874 def method(self, value: int = 1, *args, keyword=2, **kwargs):
875 try:
876 with context() as items:
877 return [item for item in items if item]
878 except Error as error:
879 match error:
880 case Error(code=code):
881 if code:
882 return f"{code!r}"
883 elif code is None:
884 return t"{code}"
885 else:
886 return "string"
887 case _:
888 return b"bytes"
889"#,
890 )
891 .expect("test source should parse");
892 let indexed = IndexedModule::new(parsed);
893 let mut visitor = Visitor {
894 nodes: Some(IndexedNodesBuilder::new()),
895 index: 0,
896 max_index: MAX_REAL_INDEX,
897 overflowed: false,
898 };
899 AnyNodeRef::from(indexed.parsed.syntax()).visit_source_order(&mut visitor);
900 let nodes = visitor
901 .nodes
902 .expect("test visitor should collect indexed nodes")
903 .all_nodes;
904
905 assert_eq!(indexed.index.len(), nodes.len());
906 let mut seen_kinds = [false; 1 << IndexedNodes::KIND_BITS];
907
908 for (raw_index, expected_node) in nodes.into_iter().enumerate() {
909 let (kind, pointer) = expected_node.into_raw_parts();
910 let address = pointer.as_ptr().expose_provenance();
911 let index = NodeIndex::from(
912 u32::try_from(raw_index).expect("node index should fit in u32"),
913 );
914 seen_kinds[usize::from(kind as u8)] = true;
915 assert_eq!(indexed.index.get(raw_index), (address, kind));
916
917 let node = indexed.get_by_index(index);
918 let (actual_kind, actual_pointer) = node.into_raw_parts();
919 assert_eq!(actual_kind, kind);
920 assert_eq!(actual_pointer.as_ptr().expose_provenance(), address);
921 assert_eq!(node.node_index().load(), index);
922 }
923 for kind in RootNodeKind::ALL {
924 let is_indexed = !matches!(
925 kind,
926 RootNodeKind::Mod | RootNodeKind::InterpolatedStringFormatSpec
927 );
928 assert_eq!(seen_kinds[usize::from(*kind as u8)], is_indexed);
929 }
930 }
931 }
932}
933
934#[cfg(test)]
935mod tests {
936 use crate::Db;
937 use crate::PythonFile;
938 use crate::files::{system_path_to_file, vendored_path_to_file};
939 use crate::parsed::parsed_module;
940 use crate::system::{
941 DbWithTestSystem, DbWithWritableSystem as _, SystemPath, SystemVirtualPath,
942 };
943 use crate::tests::TestDb;
944 use crate::vendored::{VendoredFileSystemBuilder, VendoredPath};
945 use ruff_python_ast::PythonVersion;
946 use zip::CompressionMethod;
947
948 #[test]
949 fn python_file() -> crate::system::Result<()> {
950 let mut db = TestDb::new();
951 let path = "test.py";
952
953 db.write_file(path, "x = 10")?;
954
955 let file = system_path_to_file(&db, path).unwrap();
956
957 let file = PythonFile::new(&db, file, PythonVersion::latest_ty());
958 let parsed = parsed_module(&db, file).load(&db);
959
960 assert!(parsed.has_valid_syntax());
961
962 Ok(())
963 }
964
965 #[test]
966 fn python_ipynb_file() -> crate::system::Result<()> {
967 let mut db = TestDb::new();
968 let path = SystemPath::new("test.ipynb");
969
970 db.write_file(path, "%timeit a = b")?;
971
972 let file = system_path_to_file(&db, path).unwrap();
973
974 let file = PythonFile::new(&db, file, PythonVersion::latest_ty());
975 let parsed = parsed_module(&db, file).load(&db);
976
977 assert!(parsed.has_valid_syntax());
978
979 Ok(())
980 }
981
982 #[test]
983 fn virtual_python_file() -> crate::system::Result<()> {
984 let mut db = TestDb::new();
985 let path = SystemVirtualPath::new("untitled:Untitled-1");
986
987 db.write_virtual_file(path, "x = 10");
988
989 let virtual_file = db.files().virtual_file(&db, path);
990
991 let file = PythonFile::new(&db, virtual_file.file(), PythonVersion::latest_ty());
992 let parsed = parsed_module(&db, file).load(&db);
993
994 assert!(parsed.has_valid_syntax());
995
996 Ok(())
997 }
998
999 #[test]
1000 fn virtual_ipynb_file() -> crate::system::Result<()> {
1001 let mut db = TestDb::new();
1002 let path = SystemVirtualPath::new("untitled:Untitled-1.ipynb");
1003
1004 db.write_virtual_file(path, "%timeit a = b");
1005
1006 let virtual_file = db.files().virtual_file(&db, path);
1007
1008 let file = PythonFile::new(&db, virtual_file.file(), PythonVersion::latest_ty());
1009 let parsed = parsed_module(&db, file).load(&db);
1010
1011 assert!(parsed.has_valid_syntax());
1012
1013 Ok(())
1014 }
1015
1016 #[test]
1017 fn vendored_file() {
1018 let mut db = TestDb::new();
1019
1020 let mut vendored_builder = VendoredFileSystemBuilder::new(CompressionMethod::Stored);
1021 vendored_builder
1022 .add_file(
1023 "path.pyi",
1024 r#"
1025import sys
1026
1027if sys.platform == "win32":
1028 from ntpath import *
1029 from ntpath import __all__ as __all__
1030else:
1031 from posixpath import *
1032 from posixpath import __all__ as __all__"#,
1033 )
1034 .unwrap();
1035 let vendored = vendored_builder.finish().unwrap();
1036 db.with_vendored(vendored);
1037
1038 let file = vendored_path_to_file(&db, VendoredPath::new("path.pyi")).unwrap();
1039
1040 let file = PythonFile::new(&db, file, PythonVersion::latest_ty());
1041 let parsed = parsed_module(&db, file).load(&db);
1042
1043 assert!(parsed.has_valid_syntax());
1044 }
1045
1046 #[test]
1047 fn same_file_at_different_python_versions() -> crate::system::Result<()> {
1048 let mut db = TestDb::new();
1049 db.write_file("test.py", "type Alias = int")?;
1050 let file = system_path_to_file(&db, "test.py").unwrap();
1051
1052 let py311 = PythonFile::new(&db, file, PythonVersion::PY311);
1053 let py312 = PythonFile::new(&db, file, PythonVersion::PY312);
1054 let parsed_py311 = parsed_module(&db, py311);
1055 let parsed_py312 = parsed_module(&db, py312);
1056
1057 for _ in 0..2 {
1058 assert!(
1059 !parsed_py311
1060 .load(&db)
1061 .unsupported_syntax_errors()
1062 .is_empty()
1063 );
1064 assert!(
1065 parsed_py312
1066 .load(&db)
1067 .unsupported_syntax_errors()
1068 .is_empty()
1069 );
1070
1071 parsed_py311.clear();
1072 parsed_py312.clear();
1073 }
1074
1075 Ok(())
1076 }
1077}