From d45209f425cab4b4ea34cc6080bfda8073e6b19a Mon Sep 17 00:00:00 2001 From: Andrew Gallant Date: Thu, 18 Sep 2025 09:00:21 -0400 Subject: [PATCH] [ty] Add some tricky test cases for the auto-import importer We don't attempt to fix these yet. I think there are bigger fish to fry. I came up with these based on this discussion: https://github.com/astral-sh/ruff/pull/20439#discussion_r2357769518 Here's one example: ``` if ...: from foo import MAGIC else: from bar import MAGIC MAG ``` Now in this example, completions will include `MAGIC` from the local scope. That is, auto-import is involved with that completion. But at present, auto-import will suggest importing `foo` and `bar` because we haven't de-duplicated completions yet. Which is fine. Here's another example: ``` if ...: import foo as fubar else: import bar as fubar MAG ``` Now here, there is no `MAGIC` symbol in scope. So auto-import is in play. Let's assume that the user selects `MAGIC` from `foo` in this example. (`bar` also has `MAGIC`.) Since we currently ignore the declaration site for symbols with multiple possible bindings, the importer today doesn't know that `fubar` _could_ contain `MAGIC`. But even if it did, what would we do with that information? Should we do this? ``` if ...: import foo as fubar from foo import MAGIC else: import bar as fubar MAGIC ``` Or could we reason that `bar` also has `MAGIC`? ``` if ...: import foo as fubar else: import bar as fubar fubar.MAGIC ``` But if we did that, we're making an assumption of user intent, since they *selected* `foo.MAGIC` but not `bar.MAGIC`. Anyway, I don't think we need to settle on an answer today, but I wanted to capture some of these tricky cases in tests at the very least. --- crates/ty_ide/src/importer.rs | 344 ++++++++++++++++++++++++++++++++++ 1 file changed, 344 insertions(+) diff --git a/crates/ty_ide/src/importer.rs b/crates/ty_ide/src/importer.rs index 7e698140e7..bbea834f30 100644 --- a/crates/ty_ide/src/importer.rs +++ b/crates/ty_ide/src/importer.rs @@ -1713,4 +1713,348 @@ import numpy as np (Bar) "); } + + #[test] + fn conditional_imports_new_import() { + let test = CursorTest::builder() + .source("foo.py", "MAGIC = 1") + .source("bar.py", "MAGIC = 2") + .source("quux.py", "MAGIC = 3") + .source( + "main.py", + "\ +if os.getenv(\"WHATEVER\"): + from foo import MAGIC +else: + from bar import MAGIC + +() + ", + ) + .build(); + + assert_snapshot!( + test.import("quux", "MAGIC"), @r#" + import quux + if os.getenv("WHATEVER"): + from foo import MAGIC + else: + from bar import MAGIC + + (quux.MAGIC) + "#); + assert_snapshot!( + test.import_from("quux", "MAGIC"), @r#" + import quux + if os.getenv("WHATEVER"): + from foo import MAGIC + else: + from bar import MAGIC + + (quux.MAGIC) + "#); + } + + // FIXME: This test (and the one below it) aren't + // quite right. Namely, because we aren't handling + // multiple binding sites correctly, we don't see the + // existing `MAGIC` symbol. + #[test] + fn conditional_imports_existing_import1() { + let test = CursorTest::builder() + .source("foo.py", "MAGIC = 1") + .source("bar.py", "MAGIC = 2") + .source("quux.py", "MAGIC = 3") + .source( + "main.py", + "\ +if os.getenv(\"WHATEVER\"): + from foo import MAGIC +else: + from bar import MAGIC + +() + ", + ) + .build(); + + assert_snapshot!( + test.import("foo", "MAGIC"), @r#" + import foo + if os.getenv("WHATEVER"): + from foo import MAGIC + else: + from bar import MAGIC + + (foo.MAGIC) + "#); + assert_snapshot!( + test.import_from("foo", "MAGIC"), @r#" + import foo + if os.getenv("WHATEVER"): + from foo import MAGIC + else: + from bar import MAGIC + + (foo.MAGIC) + "#); + } + + #[test] + fn conditional_imports_existing_import2() { + let test = CursorTest::builder() + .source("foo.py", "MAGIC = 1") + .source("bar.py", "MAGIC = 2") + .source("quux.py", "MAGIC = 3") + .source( + "main.py", + "\ +if os.getenv(\"WHATEVER\"): + from foo import MAGIC +else: + from bar import MAGIC + +() + ", + ) + .build(); + + assert_snapshot!( + test.import("bar", "MAGIC"), @r#" + import bar + if os.getenv("WHATEVER"): + from foo import MAGIC + else: + from bar import MAGIC + + (bar.MAGIC) + "#); + assert_snapshot!( + test.import_from("bar", "MAGIC"), @r#" + import bar + if os.getenv("WHATEVER"): + from foo import MAGIC + else: + from bar import MAGIC + + (bar.MAGIC) + "#); + } + + // FIXME: This test (and the one below it) aren't quite right. We + // don't recognize the multiple declaration sites for `fubar`. + // + // In this case, it's not totally clear what we should do. Since we + // are trying to import `MAGIC` from `foo`, we could add a `from + // foo import MAGIC` within the first `if` block. Or we could try + // and "infer" something about the code assuming that we know + // `MAGIC` is in both `foo` and `bar`. + #[test] + fn conditional_imports_existing_module1() { + let test = CursorTest::builder() + .source("foo.py", "MAGIC = 1") + .source("bar.py", "MAGIC = 2") + .source("quux.py", "MAGIC = 3") + .source( + "main.py", + "\ +if os.getenv(\"WHATEVER\"): + import foo as fubar +else: + import bar as fubar + +() + ", + ) + .build(); + + assert_snapshot!( + test.import("foo", "MAGIC"), @r#" + import foo + if os.getenv("WHATEVER"): + import foo as fubar + else: + import bar as fubar + + (foo.MAGIC) + "#); + assert_snapshot!( + test.import_from("foo", "MAGIC"), @r#" + from foo import MAGIC + if os.getenv("WHATEVER"): + import foo as fubar + else: + import bar as fubar + + (MAGIC) + "#); + } + + #[test] + fn conditional_imports_existing_module2() { + let test = CursorTest::builder() + .source("foo.py", "MAGIC = 1") + .source("bar.py", "MAGIC = 2") + .source("quux.py", "MAGIC = 3") + .source( + "main.py", + "\ +if os.getenv(\"WHATEVER\"): + import foo as fubar +else: + import bar as fubar + +() + ", + ) + .build(); + + assert_snapshot!( + test.import("bar", "MAGIC"), @r#" + import bar + if os.getenv("WHATEVER"): + import foo as fubar + else: + import bar as fubar + + (bar.MAGIC) + "#); + assert_snapshot!( + test.import_from("bar", "MAGIC"), @r#" + from bar import MAGIC + if os.getenv("WHATEVER"): + import foo as fubar + else: + import bar as fubar + + (MAGIC) + "#); + } + + #[test] + fn try_imports_new_import() { + let test = CursorTest::builder() + .source("foo.py", "MAGIC = 1") + .source("bar.py", "MAGIC = 2") + .source("quux.py", "MAGIC = 3") + .source( + "main.py", + "\ +try: + from foo import MAGIC +except ImportError: + from bar import MAGIC + +() + ", + ) + .build(); + + assert_snapshot!( + test.import("quux", "MAGIC"), @r" + import quux + try: + from foo import MAGIC + except ImportError: + from bar import MAGIC + + (quux.MAGIC) + "); + assert_snapshot!( + test.import_from("quux", "MAGIC"), @r" + import quux + try: + from foo import MAGIC + except ImportError: + from bar import MAGIC + + (quux.MAGIC) + "); + } + + // FIXME: This test (and the one below it) aren't + // quite right. Namely, because we aren't handling + // multiple binding sites correctly, we don't see the + // existing `MAGIC` symbol. + #[test] + fn try_imports_existing_import1() { + let test = CursorTest::builder() + .source("foo.py", "MAGIC = 1") + .source("bar.py", "MAGIC = 2") + .source("quux.py", "MAGIC = 3") + .source( + "main.py", + "\ +try: + from foo import MAGIC +except ImportError: + from bar import MAGIC + +() + ", + ) + .build(); + + assert_snapshot!( + test.import("foo", "MAGIC"), @r" + import foo + try: + from foo import MAGIC + except ImportError: + from bar import MAGIC + + (foo.MAGIC) + "); + assert_snapshot!( + test.import_from("foo", "MAGIC"), @r" + import foo + try: + from foo import MAGIC + except ImportError: + from bar import MAGIC + + (foo.MAGIC) + "); + } + + #[test] + fn try_imports_existing_import2() { + let test = CursorTest::builder() + .source("foo.py", "MAGIC = 1") + .source("bar.py", "MAGIC = 2") + .source("quux.py", "MAGIC = 3") + .source( + "main.py", + "\ +try: + from foo import MAGIC +except ImportError: + from bar import MAGIC + +() + ", + ) + .build(); + + assert_snapshot!( + test.import("bar", "MAGIC"), @r" + import bar + try: + from foo import MAGIC + except ImportError: + from bar import MAGIC + + (bar.MAGIC) + "); + assert_snapshot!( + test.import_from("bar", "MAGIC"), @r" + import bar + try: + from foo import MAGIC + except ImportError: + from bar import MAGIC + + (bar.MAGIC) + "); + } }