diff --git a/pyrefly/lib/stubgen.rs b/pyrefly/lib/stubgen.rs index f89da4b29c..781f308601 100644 --- a/pyrefly/lib/stubgen.rs +++ b/pyrefly/lib/stubgen.rs @@ -329,6 +329,11 @@ class BaseModel: assert_stubgen_snapshot("imports"); } + #[test] + fn test_stubgen_future_import_order() { + assert_stubgen_snapshot("future_import_order"); + } + #[test] fn test_stubgen_try_except_imports() { assert_stubgen_snapshot("try_except_imports"); diff --git a/pyrefly/lib/stubgen/emit.rs b/pyrefly/lib/stubgen/emit.rs index 2e90d6668b..56380e4d82 100644 --- a/pyrefly/lib/stubgen/emit.rs +++ b/pyrefly/lib/stubgen/emit.rs @@ -17,6 +17,21 @@ use crate::stubgen::extract::StubVariable; /// Generate the full text of a `.pyi` stub file from a `ModuleStub`. pub fn emit_stub(stub: &ModuleStub) -> String { let mut out = String::new(); + let has_future_imports = stub.items.iter().any(is_future_import); + let has_generated_imports = !stub.typing_imports.is_empty() || stub.uses_incomplete; + + for item in &stub.items { + if let StubItem::Import(imp) = item + && imp.is_future + { + out.push_str(&imp.text); + out.push('\n'); + } + } + + if has_future_imports && has_generated_imports { + out.push('\n'); + } if !stub.typing_imports.is_empty() { let names = stub.typing_imports.iter().copied().collect::>(); @@ -31,7 +46,8 @@ pub fn emit_stub(stub: &ModuleStub) -> String { out.push('\n'); } - emit_items(&stub.items, &mut out, ""); + let prev_kind = (has_future_imports && !has_generated_imports).then_some(ItemKind::Import); + emit_items(&stub.items, &mut out, "", prev_kind); if !out.ends_with('\n') { out.push('\n'); @@ -41,10 +57,12 @@ pub fn emit_stub(stub: &ModuleStub) -> String { } /// Emit a list of stub items at a given indentation level. -fn emit_items(items: &[StubItem], out: &mut String, indent: &str) { - let mut prev_kind: Option = None; - +fn emit_items(items: &[StubItem], out: &mut String, indent: &str, mut prev_kind: Option) { for item in items { + if is_future_import(item) { + continue; + } + let kind = item_kind(item); if let Some(prev) = prev_kind { @@ -80,6 +98,10 @@ fn emit_items(items: &[StubItem], out: &mut String, indent: &str) { } } +fn is_future_import(item: &StubItem) -> bool { + matches!(item, StubItem::Import(imp) if imp.is_future) +} + fn emit_function(func: &StubFunction, out: &mut String, indent: &str) { for dec in &func.decorators { out.push_str(indent); @@ -180,7 +202,7 @@ fn emit_class(cls: &StubClass, out: &mut String, indent: &str) { out.push_str(&body_indent); out.push_str("...\n"); } else { - emit_items(&cls.body, out, &body_indent); + emit_items(&cls.body, out, &body_indent, None); // If the class body only had items that were filtered out, // emit `...` as the body. diff --git a/pyrefly/lib/stubgen/extract.rs b/pyrefly/lib/stubgen/extract.rs index fd43fdfe5a..f91e501ddd 100644 --- a/pyrefly/lib/stubgen/extract.rs +++ b/pyrefly/lib/stubgen/extract.rs @@ -70,6 +70,7 @@ pub enum StubItem { pub struct StubImport { pub text: String, + pub is_future: bool, } pub struct StubFunction { @@ -197,11 +198,19 @@ fn extract_stmts( } Stmt::Import(import) => { let text = source_text(ctx.module_info, import.range()).to_owned(); - items.push(StubItem::Import(StubImport { text })); + items.push(StubItem::Import(StubImport { + text, + is_future: false, + })); } Stmt::ImportFrom(import) => { let text = source_text(ctx.module_info, import.range()).to_owned(); - items.push(StubItem::Import(StubImport { text })); + let is_future = import.level == 0 + && import + .module + .as_ref() + .is_some_and(|module| module.id == "__future__"); + items.push(StubItem::Import(StubImport { text, is_future })); } Stmt::AnnAssign(ann_assign) => { if let Some(item) = extract_ann_assign(ann_assign, ctx, in_class) { diff --git a/pyrefly/lib/test/stubgen/future_import_order/expected.pyi b/pyrefly/lib/test/stubgen/future_import_order/expected.pyi new file mode 100644 index 0000000000..d8a5fb7d81 --- /dev/null +++ b/pyrefly/lib/test/stubgen/future_import_order/expected.pyi @@ -0,0 +1,10 @@ +# @generated +from __future__ import annotations + +from typing import Literal +from _typeshed import Incomplete + +X: Literal['a'] = "a" + + +def g(v: Incomplete) -> Incomplete: ... diff --git a/pyrefly/lib/test/stubgen/future_import_order/input.py b/pyrefly/lib/test/stubgen/future_import_order/input.py new file mode 100644 index 0000000000..80ce74f860 --- /dev/null +++ b/pyrefly/lib/test/stubgen/future_import_order/input.py @@ -0,0 +1,12 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +X = "a" + + +def g(v): + return v