fix(stubgen): preserve from __future__ import annotations

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");

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::<Vec<_>>();
@@ -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<ItemKind> = None;
-
+fn emit_items(items: &[StubItem], out: &mut String, indent: &str, mut prev_kind: Option<ItemKind>) {
     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.

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) {

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: ...

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