Skip to content

Commit

Permalink
Fix/#13070 defer annotations when future is active (#13395)
Browse files Browse the repository at this point in the history
  • Loading branch information
Slyces authored Sep 19, 2024
1 parent d3530ab commit a8d9104
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 10 deletions.
9 changes: 9 additions & 0 deletions crates/red_knot_python_semantic/src/semantic_index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,9 @@ pub(crate) struct SemanticIndex<'db> {
/// Note: We should not depend on this map when analysing other files or
/// changing a file invalidates all dependents.
ast_ids: IndexVec<FileScopeId, AstIds>,

/// Flags about the global scope (code usage impacting inference)
has_future_annotations: bool,
}

impl<'db> SemanticIndex<'db> {
Expand Down Expand Up @@ -215,6 +218,12 @@ impl<'db> SemanticIndex<'db> {
pub(crate) fn node_scope(&self, node: NodeWithScopeRef) -> FileScopeId {
self.scopes_by_node[&node.node_key()]
}

/// Checks if there is an import of `__future__.annotations` in the global scope, which affects
/// the logic for type inference.
pub(super) fn has_future_annotations(&self) -> bool {
self.has_future_annotations
}
}

pub struct AncestorsIter<'a> {
Expand Down
15 changes: 15 additions & 0 deletions crates/red_knot_python_semantic/src/semantic_index/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,9 @@ pub(super) struct SemanticIndexBuilder<'db> {
/// Flow states at each `break` in the current loop.
loop_break_states: Vec<FlowSnapshot>,

/// Flags about the file's global scope
has_future_annotations: bool,

// Semantic Index fields
scopes: IndexVec<FileScopeId, Scope>,
scope_ids_by_scope: IndexVec<FileScopeId, ScopeId<'db>>,
Expand All @@ -68,6 +71,8 @@ impl<'db> SemanticIndexBuilder<'db> {
current_match_case: None,
loop_break_states: vec![],

has_future_annotations: false,

scopes: IndexVec::new(),
symbol_tables: IndexVec::new(),
ast_ids: IndexVec::new(),
Expand Down Expand Up @@ -450,6 +455,7 @@ impl<'db> SemanticIndexBuilder<'db> {
scopes_by_expression: self.scopes_by_expression,
scopes_by_node: self.scopes_by_node,
use_def_maps,
has_future_annotations: self.has_future_annotations,
}
}
}
Expand Down Expand Up @@ -543,7 +549,16 @@ where
&alias.name.id
};

// Look for imports `from __future__ import annotations`, ignore `as ...`
// We intentionally don't enforce the rules about location of `__future__`
// imports here, we assume the user's intent was to apply the `__future__`
// import, so we still check using it (and will also emit a diagnostic about a
// miss-placed `__future__` import.)
self.has_future_annotations |= alias.name.id == "annotations"
&& node.module.as_deref() == Some("__future__");

let symbol = self.add_symbol(symbol_name.clone());

self.add_definition(symbol, ImportFromDefinitionNodeRef { node, alias_index });
}
}
Expand Down
76 changes: 66 additions & 10 deletions crates/red_knot_python_semantic/src/types/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -318,9 +318,10 @@ impl<'db> TypeInferenceBuilder<'db> {
self.types.has_deferred |= inference.has_deferred;
}

/// Are we currently inferring types in a stub file?
fn is_stub(&self) -> bool {
self.file.is_stub(self.db.upcast())
/// Are we currently inferring types in file with deferred types?
/// This is true for stub files and files with `__future__.annotations`
fn are_all_types_deferred(&self) -> bool {
self.index.has_future_annotations() || self.file.is_stub(self.db.upcast())
}

/// Are we currently inferring deferred types?
Expand Down Expand Up @@ -703,7 +704,7 @@ impl<'db> TypeInferenceBuilder<'db> {
self.infer_parameters(parameters);

// TODO: this should also be applied to parameter annotations.
if self.is_stub() {
if self.are_all_types_deferred() {
self.types.has_deferred = true;
} else {
self.infer_optional_annotation_expression(returns.as_deref());
Expand Down Expand Up @@ -831,9 +832,9 @@ impl<'db> TypeInferenceBuilder<'db> {
self.infer_expression(&keyword.value);
}

// inference of bases deferred in stubs
// Inference of bases deferred in stubs
// TODO also defer stringified generic type parameters
if self.is_stub() {
if self.are_all_types_deferred() {
self.types.has_deferred = true;
} else {
for base in class.bases() {
Expand All @@ -843,13 +844,11 @@ impl<'db> TypeInferenceBuilder<'db> {
}

fn infer_function_deferred(&mut self, function: &ast::StmtFunctionDef) {
if self.is_stub() {
self.infer_optional_annotation_expression(function.returns.as_deref());
}
self.infer_optional_annotation_expression(function.returns.as_deref());
}

fn infer_class_deferred(&mut self, class: &ast::StmtClassDef) {
if self.is_stub() {
if self.are_all_types_deferred() {
for base in class.bases() {
self.infer_expression(base);
}
Expand Down Expand Up @@ -4166,6 +4165,63 @@ mod tests {
Ok(())
}

#[test]
fn deferred_annotation_in_stubs_always_resolve() -> anyhow::Result<()> {
let mut db = setup_db();

// Stub files should always resolve deferred annotations
db.write_dedented(
"/src/stub.pyi",
"
def get_foo() -> Foo: ...
class Foo: ...
foo = get_foo()
",
)?;
assert_public_ty(&db, "/src/stub.pyi", "foo", "Foo");

Ok(())
}

#[test]
fn deferred_annotations_regular_source_fails() -> anyhow::Result<()> {
let mut db = setup_db();

// In (regular) source files, deferred annotations are *not* resolved
// Also tests imports from `__future__` that are not annotations
db.write_dedented(
"/src/source.py",
"
from __future__ import with_statement as annotations
def get_foo() -> Foo: ...
class Foo: ...
foo = get_foo()
",
)?;
assert_public_ty(&db, "/src/source.py", "foo", "Unknown");

Ok(())
}

#[test]
fn deferred_annotation_in_sources_with_future_resolves() -> anyhow::Result<()> {
let mut db = setup_db();

// In source files with `__future__.annotations`, deferred annotations are resolved
db.write_dedented(
"/src/source_with_future.py",
"
from __future__ import annotations
def get_foo() -> Foo: ...
class Foo: ...
foo = get_foo()
",
)?;
assert_public_ty(&db, "/src/source_with_future.py", "foo", "Foo");

Ok(())
}

#[test]
fn narrow_not_none() -> anyhow::Result<()> {
let mut db = setup_db();
Expand Down

0 comments on commit a8d9104

Please sign in to comment.