Skip to content

Commit

Permalink
test helper
Browse files Browse the repository at this point in the history
  • Loading branch information
carljm committed Aug 16, 2024
1 parent 5d70f79 commit ad6863d
Showing 1 changed file with 38 additions and 49 deletions.
87 changes: 38 additions & 49 deletions crates/red_knot_python_semantic/src/semantic_index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -311,13 +311,29 @@ mod tests {
use ruff_text_size::{Ranged, TextRange};

use crate::db::tests::TestDb;
use crate::semantic_index::ast_ids::HasScopedUseId;
use crate::semantic_index::definition::DefinitionKind;
use crate::semantic_index::symbol::{FileScopeId, Scope, ScopeKind, SymbolTable};
use crate::semantic_index::use_def::DefinitionWithConstraints;
use crate::semantic_index::ast_ids::{HasScopedUseId, ScopedUseId};
use crate::semantic_index::definition::{Definition, DefinitionKind};
use crate::semantic_index::symbol::{
FileScopeId, Scope, ScopeKind, ScopedSymbolId, SymbolTable,
};
use crate::semantic_index::use_def::UseDefMap;
use crate::semantic_index::{global_scope, semantic_index, symbol_table, use_def_map};
use crate::Db;

impl UseDefMap<'_> {
fn first_public_definition(&self, symbol: ScopedSymbolId) -> Option<Definition<'_>> {
self.public_definitions(symbol)
.next()
.map(|constrained_definition| constrained_definition.definition)
}

fn first_use_definition(&self, use_id: ScopedUseId) -> Option<Definition<'_>> {
self.use_definitions(use_id)
.next()
.map(|constrained_definition| constrained_definition.definition)
}
}

struct TestCase {
db: TestDb,
file: File,
Expand Down Expand Up @@ -376,11 +392,7 @@ mod tests {
let foo = global_table.symbol_id_by_name("foo").unwrap();

let use_def = use_def_map(&db, scope);
let [DefinitionWithConstraints { definition, .. }] =
use_def.public_definitions(foo).collect::<Vec<_>>()[..]
else {
panic!("expected one definition");
};
let definition = use_def.first_public_definition(foo).unwrap();
assert!(matches!(definition.node(&db), DefinitionKind::Import(_)));
}

Expand Down Expand Up @@ -415,16 +427,13 @@ mod tests {
);

let use_def = use_def_map(&db, scope);
let [DefinitionWithConstraints { definition, .. }] = use_def
.public_definitions(
let definition = use_def
.first_public_definition(
global_table
.symbol_id_by_name("foo")
.expect("symbol to exist"),
)
.collect::<Vec<_>>()[..]
else {
panic!("expected one definition");
};
.unwrap();
assert!(matches!(
definition.node(&db),
DefinitionKind::ImportFrom(_)
Expand All @@ -445,12 +454,9 @@ mod tests {
"a symbol used but not defined in a scope should have only the used flag"
);
let use_def = use_def_map(&db, scope);
let [DefinitionWithConstraints { definition, .. }] = use_def
.public_definitions(global_table.symbol_id_by_name("x").expect("symbol exists"))
.collect::<Vec<_>>()[..]
else {
panic!("expected one definition");
};
let definition = use_def
.first_public_definition(global_table.symbol_id_by_name("x").expect("symbol exists"))
.unwrap();
assert!(matches!(
definition.node(&db),
DefinitionKind::Assignment(_)
Expand Down Expand Up @@ -485,12 +491,9 @@ y = 2
assert_eq!(names(&class_table), vec!["x"]);

let use_def = index.use_def_map(class_scope_id);
let [DefinitionWithConstraints { definition, .. }] = use_def
.public_definitions(class_table.symbol_id_by_name("x").expect("symbol exists"))
.collect::<Vec<_>>()[..]
else {
panic!("expected one definition");
};
let definition = use_def
.first_public_definition(class_table.symbol_id_by_name("x").expect("symbol exists"))
.unwrap();
assert!(matches!(
definition.node(&db),
DefinitionKind::Assignment(_)
Expand Down Expand Up @@ -524,16 +527,13 @@ y = 2
assert_eq!(names(&function_table), vec!["x"]);

let use_def = index.use_def_map(function_scope_id);
let [DefinitionWithConstraints { definition, .. }] = use_def
.public_definitions(
let definition = use_def
.first_public_definition(
function_table
.symbol_id_by_name("x")
.expect("symbol exists"),
)
.collect::<Vec<_>>()[..]
else {
panic!("expected one definition");
};
.unwrap();
assert!(matches!(
definition.node(&db),
DefinitionKind::Assignment(_)
Expand Down Expand Up @@ -606,11 +606,7 @@ y = 2
let element_use_id =
element.scoped_use_id(&db, comprehension_scope_id.to_scope_id(&db, file));

let [DefinitionWithConstraints { definition, .. }] =
use_def.use_definitions(element_use_id).collect::<Vec<_>>()[..]
else {
panic!("expected one definition")
};
let definition = use_def.first_use_definition(element_use_id).unwrap();
let DefinitionKind::Comprehension(comprehension) = definition.node(&db) else {
panic!("expected generator definition")
};
Expand Down Expand Up @@ -707,16 +703,13 @@ def func():
assert_eq!(names(&func2_table), vec!["y"]);

let use_def = index.use_def_map(FileScopeId::global());
let [DefinitionWithConstraints { definition, .. }] = use_def
.public_definitions(
let definition = use_def
.first_public_definition(
global_table
.symbol_id_by_name("func")
.expect("symbol exists"),
)
.collect::<Vec<_>>()[..]
else {
panic!("expected one definition");
};
.unwrap();
assert!(matches!(definition.node(&db), DefinitionKind::Function(_)));
}

Expand Down Expand Up @@ -817,11 +810,7 @@ class C[T]:
};
let x_use_id = x_use_expr_name.scoped_use_id(&db, scope);
let use_def = use_def_map(&db, scope);
let [DefinitionWithConstraints { definition, .. }] =
use_def.use_definitions(x_use_id).collect::<Vec<_>>()[..]
else {
panic!("expected one definition");
};
let definition = use_def.first_use_definition(x_use_id).unwrap();
let DefinitionKind::Assignment(assignment) = definition.node(&db) else {
panic!("should be an assignment definition")
};
Expand Down

0 comments on commit ad6863d

Please sign in to comment.