[ty] Include conditional symbols (like datetime.UTC) in auto-import in more cases

This is a quick fix to make conditional symbols in auto-import work. In
essence, this flips the current failure mode of "don't suggest a symbol
that is available" to "possibly suggest a symbol that isn't available."
I think suggesting a symbol that isn't available is probably the better
failure mode.

Fixes #1758, Ref #2795

Author: BurntSushi

crates/ty_ide/src/symbols.rs

@@ -1101,12 +1101,30 @@ impl<'db> SymbolVisitor<'db> {
     }
 
     /// Updates the origin of `__all__` in the current module.
-    ///
-    /// This will clear existing names if the origin is changed to
-    /// mimic the behavior of overriding `__all__` in the current
-    /// module.
     fn update_all_origin(&mut self, origin: DunderAllOrigin) {
-        if self.all_origin.is_some() {
+        // N.B. This used to clear `all_names` whenever there
+        // was *any* previous origin set. Now we skip clearing
+        // it if the previous origin and the new origin are
+        // both "current module." This tends to arise in situation
+        // like this:
+        //
+        //     if sys.version > ...:
+        //         __all__ = ['SomeFancyNewSymbol']
+        //     else:
+        //         __all__ = []
+        //
+        // Clearing is arguably correct here, but auto-import
+        // (unlike ty's own __all__ handling) doesn't yet know
+        // how to evaluate conditionals. So instead of over-writing
+        // __all__, we union it. This will produce incorrect
+        // results in some cases, but the failure mode will be
+        // "suggests symbol that doesn't exist" instead of
+        // "doesn't suggest symbol that does exist." The former
+        // seems preferable (until we know how to evaluate at
+        // least some rudimentary conditionals).
+        if !(matches!(self.all_origin, Some(DunderAllOrigin::CurrentModule))
+            && matches!(origin, DunderAllOrigin::CurrentModule))
+        {
             self.all_names.clear();
         }
         self.all_origin = Some(origin);
@@ -1448,6 +1466,7 @@ mod tests {
     use ruff_db::Db;
     use ruff_db::files::{FileRootKind, system_path_to_file};
     use ruff_db::system::{DbWithWritableSystem, SystemPath, SystemPathBuf};
+    use ruff_python_ast::PythonVersion;
     use ruff_python_trivia::textwrap::dedent;
     use ty_project::{ProjectMetadata, TestDb};
 
@@ -2713,6 +2732,79 @@ class X:
         );
     }
 
+    /// Tests that a work-around which unions `__all__` values lets
+    /// us find conditionally exported symbols.
+    ///
+    /// However, this also means that we may suggest exported symbols
+    /// even when they aren't available. What we should ideally do is
+    /// evaluate the conditional like ty does, but this requires some
+    /// work. (And it's unlikely auto-import will ever get the full
+    /// evaluation capabilities as ty, so it's likely this sort of
+    /// union work-around will always be something we do in at least
+    /// some cases.)
+    #[test]
+    fn union_all_to_work_around_conditional_symbols_py311() {
+        let test = PublicTestBuilder::default()
+            .python_version(PythonVersion::PY311)
+            .source(
+                "test.py",
+                "
+                import sys
+                if sys.version_info >= (3, 11):
+                    ZQZQZQ = 1
+                else:
+                    ZYZYZY = 1
+
+                if sys.version_info >= (3, 11):
+                    __all__ = ['ZQZQZQ']
+                else:
+                    __all__ = ['ZYZYZY']
+                ",
+            )
+            .build();
+        // Ideally this would only have `ZQZQZQ`.
+        insta::assert_snapshot!(
+            test.exports(),
+            @r"
+        ZQZQZQ :: Constant
+        ZYZYZY :: Constant
+        ",
+        );
+    }
+
+    /// Like `union_all_to_work_around_conditional_symbols_py311`, but
+    /// sets the environment Python version to 3.10 so that the conditional
+    /// should evaluate to false.
+    #[test]
+    fn union_all_to_work_around_conditional_symbols_py310() {
+        let test = PublicTestBuilder::default()
+            .python_version(PythonVersion::PY310)
+            .source(
+                "test.py",
+                "
+                import sys
+                if sys.version_info >= (3, 11):
+                    ZQZQZQ = 1
+                else:
+                    ZYZYZY = 1
+
+                if sys.version_info >= (3, 11):
+                    __all__ = ['ZQZQZQ']
+                else:
+                    __all__ = ['ZYZYZY']
+                ",
+            )
+            .build();
+        // Ideally this would only have `ZYZYZY`.
+        insta::assert_snapshot!(
+            test.exports(),
+            @r"
+        ZQZQZQ :: Constant
+        ZYZYZY :: Constant
+        ",
+        );
+    }
+
     #[test]
     fn deprecated_function() {
         let test = public_test(
@@ -2848,16 +2940,17 @@ class C: ...
         /// A list of source files, corresponding to the
         /// file's path and its contents.
         sources: Vec<Source>,
+        /// The python version to use.
+        python_version: Option<PythonVersion>,
     }
 
     impl PublicTestBuilder {
         pub(super) fn build(&self) -> PublicTest {
-            let mut db = TestDb::new(ProjectMetadata::new(
-                "test".into(),
-                SystemPathBuf::from("/"),
-            ));
+            let metadata = ProjectMetadata::new("test".into(), SystemPathBuf::from("/"));
+            let mut db = TestDb::new(metadata);
 
-            db.init_program().unwrap();
+            db.init_program_with_python_version(self.python_version.unwrap_or_default())
+                .unwrap();
 
             for Source { path, contents } in &self.sources {
                 db.write_file(path, contents)
@@ -2897,6 +2990,11 @@ class C: ...
             self.sources.push(Source { path, contents });
             self
         }
+
+        pub(super) fn python_version(&mut self, version: PythonVersion) -> &mut PublicTestBuilder {
+            self.python_version = Some(version);
+            self
+        }
     }
 
     struct Source {

crates/ty_project/src/db.rs

@@ -542,6 +542,7 @@ pub(crate) mod tests {
     use ruff_db::files::{FileRootKind, Files};
     use ruff_db::system::{DbWithTestSystem, System, TestSystem};
     use ruff_db::vendored::VendoredFileSystem;
+    use ruff_python_ast::PythonVersion;
     use ty_module_resolver::SearchPathSettings;
     use ty_python_semantic::lint::{LintRegistry, RuleSelection};
     use ty_python_semantic::{
@@ -588,6 +589,13 @@ pub(crate) mod tests {
         }
 
         pub fn init_program(&mut self) -> anyhow::Result<()> {
+            self.init_program_with_python_version(PythonVersion::latest_ty())
+        }
+
+        pub fn init_program_with_python_version(
+            &mut self,
+            python_version: PythonVersion,
+        ) -> anyhow::Result<()> {
             let root = self.project().root(self);
 
             let search_paths = SearchPathSettings::new(vec![root.to_path_buf()])
@@ -597,7 +605,10 @@ pub(crate) mod tests {
             Program::from_settings(
                 self,
                 ProgramSettings {
-                    python_version: PythonVersionWithSource::default(),
+                    python_version: PythonVersionWithSource {
+                        source: ty_python_semantic::PythonVersionSource::Default,
+                        version: python_version,
+                    },
                     python_platform: PythonPlatform::default(),
                     search_paths,
                 },