Consider from thispackage import y to re-export y in __init__.pyi

Fixes https://github.com/astral-sh/ty/issues/1487

crates/ty_python_semantic/resources/mdtest/import/nonstandard_conventions.md

@@ -29,7 +29,8 @@ defining symbols *at all* and re-exporting them.
 
 ## Relative `from` Import of Direct Submodule in `__init__`
 
-We consider the `from . import submodule` idiom in an `__init__.pyi` an explicit re-export.
+We consider the `from . import submodule` idiom in an `__init__.pyi` an explicit re-export. This
+pattern is observed in the wild with various stub packages.
 
 ### In Stub
 
@@ -94,8 +95,7 @@ reveal_type(mypackage.fails.Y)  # revealed: Unknown
 ## Absolute `from` Import of Direct Submodule in `__init__`
 
 If an absolute `from...import` happens to import a submodule (i.e. it's equivalent to
-`from . import y`) we do not treat it as a re-export. We could, but we don't. (This is an arbitrary
-decision and can be changed!)
+`from . import y`) we also treat it as a re-export.
 
 ### In Stub
 
@@ -122,9 +122,7 @@ Y: int = 47
 ```py
 import mypackage
 
-# TODO: this could work and would be nice to have?
-# error: "has no member `imported`"
-reveal_type(mypackage.imported.X)  # revealed: Unknown
+reveal_type(mypackage.imported.X)  # revealed: int
 # error: "has no member `fails`"
 reveal_type(mypackage.fails.Y)  # revealed: Unknown
 ```
@@ -333,7 +331,7 @@ reveal_type(mypackage.nested.X)  # revealed: Unknown
 
 ### In Non-Stub
 
-`from mypackage.submodule import nested` in an `__init__.py` only creates `nested`.
+`from mypackage.submodule import nested` in an `__init__.py` creates both `submodule` and `nested`.
 
 `mypackage/__init__.py`:
 
@@ -357,12 +355,11 @@ X: int = 42
 ```py
 import mypackage
 
+reveal_type(mypackage.submodule)  # revealed: <module 'mypackage.submodule'>
 # TODO: this would be nice to support
-# error: "has no member `submodule`"
-reveal_type(mypackage.submodule)  # revealed: Unknown
-# error: "has no member `submodule`"
+# error: "has no member `nested`"
 reveal_type(mypackage.submodule.nested)  # revealed: Unknown
-# error: "has no member `submodule`"
+# error: "has no member `nested`"
 reveal_type(mypackage.submodule.nested.X)  # revealed: Unknown
 reveal_type(mypackage.nested)  # revealed: <module 'mypackage.submodule.nested'>
 reveal_type(mypackage.nested.X)  # revealed: int

crates/ty_python_semantic/src/module_name.rs

@@ -295,6 +295,7 @@ impl ModuleName {
         Self::from_identifier_parts(db, importing_file, module.as_deref(), *level)
     }
 
+    /// Computes the absolute module name from the LHS components of `from LHS import RHS`
     pub(crate) fn from_identifier_parts(
         db: &dyn Db,
         importing_file: File,
@@ -309,6 +310,16 @@ impl ModuleName {
                 .ok_or(ModuleNameResolutionError::InvalidSyntax)
         }
     }
+
+    /// Computes the absolute module name for the package this file belongs to.
+    ///
+    /// i.e. this resolves `.`
+    pub(crate) fn package_for_file(
+        db: &dyn Db,
+        importing_file: File,
+    ) -> Result<Self, ModuleNameResolutionError> {
+        Self::from_identifier_parts(db, importing_file, None, 1)
+    }
 }
 
 impl Deref for ModuleName {

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -1451,7 +1451,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
 
                 // If we see:
                 //
-                // * `from .x.y import z` (must be relative!)
+                // * `from .x.y import z` (or `from whatever.thispackage.x.y`)
                 // * And we are in an `__init__.py(i)` (hereafter `thispackage`)
                 // * And this is the first time we've seen `from .x` in this module
                 // * And we're in the global scope
@@ -1465,25 +1465,35 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 // reasons but it works well for most practical purposes. In particular it's nice
                 // that `x` can be freely overwritten, and that we don't assume that an import
                 // in one function is visible in another function.
-                //
-                // TODO: Also support `from thispackage.x.y import z`?
-                if self.current_scope() == FileScopeId::global()
-                    && node.level == 1
-                    && let Some(submodule) = &node.module
-                    && let Some(parsed_submodule) = ModuleName::new(submodule.as_str())
-                    && let Some(direct_submodule) = parsed_submodule.components().next()
-                    && self.file.is_package(self.db)
-                    && !self.seen_submodule_imports.contains(direct_submodule)
+                let mut is_self_import = false;
+                if self.file.is_package(self.db)
+                    && let Ok(module_name) = ModuleName::from_identifier_parts(
+                        self.db,
+                        self.file,
+                        node.module.as_deref(),
+                        node.level,
+                    )
+                    && let Ok(thispackage) = ModuleName::package_for_file(self.db, self.file)
                 {
-                    self.seen_submodule_imports
-                        .insert(direct_submodule.to_owned());
+                    // Record whether this is equivalent to `from . import ...`
+                    is_self_import = module_name == thispackage;
 
-                    let direct_submodule_name = Name::new(direct_submodule);
-                    let symbol = self.add_symbol(direct_submodule_name);
-                    self.add_definition(
-                        symbol.into(),
-                        ImportFromSubmoduleDefinitionNodeRef { node, submodule },
-                    );
+                    if node.module.is_some()
+                        && let Some(relative_submodule) = module_name.relative_to(&thispackage)
+                        && let Some(direct_submodule) = relative_submodule.components().next()
+                        && !self.seen_submodule_imports.contains(direct_submodule)
+                        && self.current_scope().is_global()
+                    {
+                        self.seen_submodule_imports
+                            .insert(direct_submodule.to_owned());
+
+                        let direct_submodule_name = Name::new(direct_submodule);
+                        let symbol = self.add_symbol(direct_submodule_name);
+                        self.add_definition(
+                            symbol.into(),
+                            ImportFromSubmoduleDefinitionNodeRef { node },
+                        );
+                    }
                 }
 
                 let mut found_star = false;
@@ -1595,13 +1605,10 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                         // It's re-exported if it's `from ... import x as x`
                         (&asname.id, asname.id == alias.name.id)
                     } else {
-                        // It's re-exported if it's `from . import x` in an `__init__.pyi`
-                        (
-                            &alias.name.id,
-                            node.level == 1
-                                && node.module.is_none()
-                                && self.file.is_package(self.db),
-                        )
+                        // As a non-standard rule to handle stubs in the wild, we consider
+                        // `from . import x` and `from whatever.thispackage import x` in an
+                        // `__init__.pyi` to re-export `x` (as long as it wasn't renamed)
+                        (&alias.name.id, is_self_import)
                     };
 
                     // Look for imports `from __future__ import annotations`, ignore `as ...`

crates/ty_python_semantic/src/semantic_index/definition.rs

@@ -3,7 +3,6 @@ use std::ops::Deref;
 use ruff_db::files::{File, FileRange};
 use ruff_db::parsed::{ParsedModuleRef, parsed_module};
 use ruff_python_ast as ast;
-use ruff_python_ast::name::Name;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::Db;
@@ -368,7 +367,6 @@ pub(crate) struct ImportFromDefinitionNodeRef<'ast> {
 #[derive(Copy, Clone, Debug)]
 pub(crate) struct ImportFromSubmoduleDefinitionNodeRef<'ast> {
     pub(crate) node: &'ast ast::StmtImportFrom,
-    pub(crate) submodule: &'ast ast::Identifier,
 }
 #[derive(Copy, Clone, Debug)]
 pub(crate) struct AssignmentDefinitionNodeRef<'ast, 'db> {
@@ -450,10 +448,8 @@ impl<'db> DefinitionNodeRef<'_, 'db> {
             }),
             DefinitionNodeRef::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef {
                 node,
-                submodule,
             }) => DefinitionKind::ImportFromSubmodule(ImportFromSubmoduleDefinitionKind {
                 node: AstNodeRef::new(parsed, node),
-                submodule: submodule.as_str().into(),
             }),
             DefinitionNodeRef::ImportStar(star_import) => {
                 let StarImportDefinitionNodeRef { node, symbol_id } = star_import;
@@ -580,10 +576,7 @@ impl<'db> DefinitionNodeRef<'_, 'db> {
                 alias_index,
                 is_reexported: _,
             }) => (&node.names[alias_index]).into(),
-            Self::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef {
-                node,
-                submodule: _,
-            }) => node.into(),
+            Self::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef { node }) => node.into(),
             // INVARIANT: for an invalid-syntax statement such as `from foo import *, bar, *`,
             // we only create a `StarImportDefinitionKind` for the *first* `*` alias in the names list.
             Self::ImportStar(StarImportDefinitionNodeRef { node, symbol_id: _ }) => node
@@ -1021,17 +1014,12 @@ impl ImportFromDefinitionKind {
 #[derive(Clone, Debug, get_size2::GetSize)]
 pub struct ImportFromSubmoduleDefinitionKind {
     node: AstNodeRef<ast::StmtImportFrom>,
-    submodule: Name,
 }
 
 impl ImportFromSubmoduleDefinitionKind {
     pub fn import<'ast>(&self, module: &'ast ParsedModuleRef) -> &'ast ast::StmtImportFrom {
         self.node.node(module)
     }
-
-    pub(crate) fn submodule(&self) -> &Name {
-        &self.submodule
-    }
 }
 
 #[derive(Clone, Debug, get_size2::GetSize)]

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -4,7 +4,6 @@ use itertools::{Either, Itertools};
 use ruff_db::diagnostic::{Annotation, DiagnosticId, Severity};
 use ruff_db::files::File;
 use ruff_db::parsed::ParsedModuleRef;
-use ruff_python_ast::name::Name;
 use ruff_python_ast::visitor::{Visitor, walk_expr};
 use ruff_python_ast::{
     self as ast, AnyNodeRef, ExprContext, HasNodeIndex, NodeIndex, PythonVersion,
@@ -1218,7 +1217,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             DefinitionKind::ImportFromSubmodule(import_from) => {
                 self.infer_import_from_submodule_definition(
                     import_from.import(self.module()),
-                    import_from.submodule(),
                     definition,
                 );
             }
@@ -5901,51 +5899,64 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         }
     }
 
-    /// Infer the implicit local definition `x = <module 'thispackage.x'>` that
-    /// `from .x.y import z` can introduce in an `__init__.py(i)`.
+    /// Infer the implicit local definition `x = <module 'whatever.thispackage.x'>` that
+    /// `from .x.y import z` or `from whatever.thispackage.x.y` can introduce in `__init__.py(i)`.
     ///
     /// For the definition `z`, see [`TypeInferenceBuilder::infer_import_from_definition`].
+    ///
+    /// The runtime semantic of this kind of statement is to introduce a variable in the global
+    /// scope of this module *the first time it's imported in the entire program*. This
+    /// implementation just blindly introduces a local variable wherever the `from..import` is
+    /// (if the imports actually resolve).
+    ///
+    /// That gap between the semantics and implementation are currently the responsibility of the
+    /// code that actually creates these kinds of Definitions (so blindly introducing a local
+    /// is all we need to be doing here).
     fn infer_import_from_submodule_definition(
         &mut self,
         import_from: &ast::StmtImportFrom,
-        submodule: &Name,
         definition: Definition<'db>,
     ) {
-        // The runtime semantic of this kind of statement is to introduce a variable in the global
-        // scope of this module, so we do just that. (Actually we introduce a local variable, but
-        // this type of Definition is only created when a `from..import` is in global scope.)
-
-        // Get this package's module by resolving `.`
-        let Ok(module_name) = ModuleName::from_identifier_parts(self.db(), self.file(), None, 1)
-        else {
+        // Get this package's absolute module name by resolving `.`, and make sure it exists
+        let Ok(thispackage_name) = ModuleName::package_for_file(self.db(), self.file()) else {
+            self.add_binding(import_from.into(), definition, |_, _| Type::unknown());
+            return;
+        };
+        let Some(module) = resolve_module(self.db(), &thispackage_name) else {
             self.add_binding(import_from.into(), definition, |_, _| Type::unknown());
             return;
         };
 
-        let Some(module) = resolve_module(self.db(), &module_name) else {
+        // We have `from whatever.thispackage.x.y ...` or `from .x.y ...`
+        // and we want to extract `x` (to ultimately construct `whatever.thispackage.x`):
+
+        // First we normalize to `whatever.thispackage.x.y`
+        let Some(final_part) = ModuleName::from_identifier_parts(
+            self.db(),
+            self.file(),
+            import_from.module.as_deref(),
+            import_from.level,
+        )
+        .ok()
+        // `whatever.thispackage.x.y` => `x.y`
+        .and_then(|submodule_name| submodule_name.relative_to(&thispackage_name))
+        // `x.y` => `x`
+        .and_then(|relative_submodule_name| {
+            relative_submodule_name
+                .components()
+                .next()
+                .and_then(ModuleName::new)
+        }) else {
             self.add_binding(import_from.into(), definition, |_, _| Type::unknown());
             return;
         };
 
-        // Now construct the submodule `.x`
-        assert!(
-            !submodule.is_empty(),
-            "ImportFromSubmoduleDefinitionKind constructed with empty module"
-        );
-        let name = submodule
-            .split_once('.')
-            .map(|(first, _)| first)
-            .unwrap_or(submodule.as_str());
-        let full_submodule_name = ModuleName::new(name).map(|final_part| {
-            let mut ret = module_name.clone();
-            ret.extend(&final_part);
-            ret
-        });
-        // And try to import it
-        if let Some(submodule_type) = full_submodule_name
-            .as_ref()
-            .and_then(|submodule_name| self.module_type_from_name(submodule_name))
-        {
+        // `x` => `whatever.thispackage.x`
+        let mut full_submodule_name = thispackage_name.clone();
+        full_submodule_name.extend(&final_part);
+
+        // Try to actually resolve the import `whatever.thispackage.x`
+        if let Some(submodule_type) = self.module_type_from_name(&full_submodule_name) {
             // Success, introduce a binding!
             //
             // We explicitly don't introduce a *declaration* because it's actual ok
@@ -5970,17 +5981,15 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         };
 
         let diagnostic = builder.into_diagnostic(format_args!(
-            "Module `{module_name}` has no submodule `{name}`"
+            "Module `{thispackage_name}` has no submodule `{final_part}`"
         ));
 
-        if let Some(full_submodule_name) = full_submodule_name {
-            hint_if_stdlib_submodule_exists_on_other_versions(
-                self.db(),
-                diagnostic,
-                &full_submodule_name,
-                module,
-            );
-        }
+        hint_if_stdlib_submodule_exists_on_other_versions(
+            self.db(),
+            diagnostic,
+            &full_submodule_name,
+            module,
+        );
     }
 
     fn infer_return_statement(&mut self, ret: &ast::StmtReturn) {