Repository: astral-sh/ruff
File: crates/ty_ide/src/semantic_tokens.rs (Rust)

Code before:
    #[test]
    fn import_as() {
        let test = SemanticTokenTest::new(
            r#"
            import pathlib as path
            from pathlib import Path
            "#,
        );

        let tokens = test.highlight_file();
        assert_snapshot!(test.to_snapshot(&tokens), @r#"
        "pathlib" @ 8..15: Namespace
        "path" @ 19..23: Namespace
        "pathlib" @ 29..36: Namespace
        "Path" @ 44..48: Class
        "#);
    }

    #[test]
    fn import_from_as() {
        // Test that both the imported name and its alias get highlighted
        // See: https://github.com/astral-sh/ty/issues/2547
        let test = SemanticTokenTest::new(
            r#"
from pathlib import Path as P
from typing import Set as AbstractSet
"#,
        );

        let tokens = test.highlight_file();
        // Both the imported name and the alias should get the same highlighting
        assert_snapshot!(test.to_snapshot(&tokens), @r#"
        "pathlib" @ 6..13: Namespace
        "Path" @ 21..25: Class
        "P" @ 29..30: Class
        "typing" @ 36..42: Namespace
        "Set" @ 50..53: Variable
        "AbstractSet" @ 57..68: Variable
        "#);
    }

    pub(super) struct SemanticTokenTest {
        pub(super) db: ty_project::TestDb,
        file: File,
    }

    impl SemanticTokenTest {
        fn new(source: &str) -> Self {
            let mut db = ty_project::TestDb::new(ProjectMetadata::new(
                "test".into(),
                SystemPathBuf::from("/"),

Review comment (refactor):
nittiest of nitpicks: I think you meant

```suggestion
from collections.abc import Set as AbstractSet
```

assuming that this is a reference to the convention enforced by https://docs.astral.sh/ruff/rules/unaliased-collections-abc-set-import/ 😄

Code after:
    #[test]
    fn import_as() {
        let test = SemanticTokenTest::new(
            r#"
            import pathlib as path
            from pathlib import Path
            "#,
        );

        let tokens = test.highlight_file();
        assert_snapshot!(test.to_snapshot(&tokens), @r#"
        "pathlib" @ 8..15: Namespace
        "path" @ 19..23: Namespace
        "pathlib" @ 29..36: Namespace
        "Path" @ 44..48: Class
        "#);
    }

    #[test]
    fn import_from_as() {
        // Test that both the imported name and its alias get highlighted
        // See: https://github.com/astral-sh/ty/issues/2547
        let test = SemanticTokenTest::new(
            r#"
from pathlib import Path as P
from collections.abc import Set as AbstractSet
"#,
        );

        let tokens = test.highlight_file();
        // Both the imported name and the alias should get the same highlighting
        assert_snapshot!(test.to_snapshot(&tokens), @r#"
        "pathlib" @ 6..13: Namespace
        "Path" @ 21..25: Class
        "P" @ 29..30: Class
        "collections" @ 36..47: Namespace
        "abc" @ 48..51: Namespace
        "Set" @ 59..62: Class
        "AbstractSet" @ 66..77: Class
        "#);
    }

    pub(super) struct SemanticTokenTest {
        pub(super) db: ty_project::TestDb,
        file: File,
    }

    impl SemanticTokenTest {
        fn new(source: &str) -> Self {
            let mut db = ty_project::TestDb::new(ProjectMetadata::new(
                "test".into(),