Skip to content

Commit 5b4e753

Browse files
MatthewMckee4MatthewMckee4
authored
[ty] Add support for goto in literal enum member inlay hint (#24792)
1 parent e7cc762 commit 5b4e753

6 files changed

Lines changed: 140 additions & 18 deletions

File tree

‎crates/ty_ide/src/goto.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,8 @@ impl<'db> Definitions<'db> {
257257
| ty_python_semantic::types::TypeDefinition::TypeVar(definition)
258258
| ty_python_semantic::types::TypeDefinition::TypeAlias(definition)
259259
| ty_python_semantic::types::TypeDefinition::SpecialForm(definition)
260-
| ty_python_semantic::types::TypeDefinition::NewType(definition) => {
260+
| ty_python_semantic::types::TypeDefinition::NewType(definition)
261+
| ty_python_semantic::types::TypeDefinition::EnumMember(definition) => {
261262
ResolvedDefinition::Definition(definition)
262263
}
263264
};

‎crates/ty_ide/src/inlay_hints.rs‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2777,6 +2777,86 @@ Source with applied edits:
27772777
"#);
27782778
}
27792779

2780+
#[test]
2781+
fn test_enum_literal() {
2782+
let mut test = inlay_hint_test(
2783+
r#"
2784+
from enum import Enum
2785+
2786+
class Color(Enum):
2787+
RED = 1
2788+
BLUE = 2
2789+
2790+
x = Color.RED
2791+
"#,
2792+
);
2793+
2794+
assert_snapshot!(test.inlay_hints(), @r"
2795+
2796+
from enum import Enum
2797+
2798+
class Color(Enum):
2799+
RED = 1
2800+
BLUE = 2
2801+
2802+
x[: Literal[Color.RED]] = Color.RED
2803+
2804+
---------------------------------------------
2805+
info[inlay-hint-location]: Inlay Hint Target
2806+
--> stdlib/typing.pyi:487:1
2807+
|
2808+
487 | Literal: _SpecialForm
2809+
| ^^^^^^^
2810+
|
2811+
info: Source
2812+
--> main2.py:8:5
2813+
|
2814+
8 | x[: Literal[Color.RED]] = Color.RED
2815+
| ^^^^^^^
2816+
|
2817+
2818+
info[inlay-hint-location]: Inlay Hint Target
2819+
--> main.py:4:7
2820+
|
2821+
4 | class Color(Enum):
2822+
| ^^^^^
2823+
|
2824+
info: Source
2825+
--> main2.py:8:13
2826+
|
2827+
8 | x[: Literal[Color.RED]] = Color.RED
2828+
| ^^^^^
2829+
|
2830+
2831+
info[inlay-hint-location]: Inlay Hint Target
2832+
--> main.py:5:5
2833+
|
2834+
5 | RED = 1
2835+
| ^^^
2836+
|
2837+
info: Source
2838+
--> main2.py:8:19
2839+
|
2840+
8 | x[: Literal[Color.RED]] = Color.RED
2841+
| ^^^
2842+
|
2843+
2844+
---------------------------------------------
2845+
info[inlay-hint-edit]: Inlay hint edits
2846+
--> main.py:1:1
2847+
1 + from typing import Literal
2848+
2 |
2849+
3 | from enum import Enum
2850+
4 |
2851+
--------------------------------------------------------------------------------
2852+
6 | RED = 1
2853+
7 | BLUE = 2
2854+
8 |
2855+
- x = Color.RED
2856+
9 + x: Literal[Color.RED] = Color.RED
2857+
");
2858+
}
2859+
27802860
#[test]
27812861
fn test_simple_init_call() {
27822862
let mut test = inlay_hint_test(

‎crates/ty_python_semantic/src/types.rs‎

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -6202,16 +6202,18 @@ impl<'db> Type<'db> {
62026202
KnownInstanceType::TypeAliasType(type_alias) => {
62036203
Some(TypeDefinition::TypeAlias(type_alias.definition(db)))
62046204
}
6205-
KnownInstanceType::NewType(newtype) => Some(TypeDefinition::NewType(newtype.definition(db))),
6205+
KnownInstanceType::NewType(newtype) => {
6206+
Some(TypeDefinition::NewType(newtype.definition(db)))
6207+
}
62066208
_ => None,
62076209
},
62086210

62096211
Self::SubclassOf(subclass_of_type) => match subclass_of_type.subclass_of() {
62106212
SubclassOfInner::Dynamic(_) => None,
62116213
SubclassOfInner::Class(class) => class.type_definition(db),
6212-
SubclassOfInner::TypeVar(bound_typevar) => {
6213-
Some(TypeDefinition::TypeVar(bound_typevar.typevar(db).definition(db)?))
6214-
}
6214+
SubclassOfInner::TypeVar(bound_typevar) => Some(TypeDefinition::TypeVar(
6215+
bound_typevar.typevar(db).definition(db)?,
6216+
)),
62156217
},
62166218

62176219
Self::TypeAlias(alias) => alias.value_type(db).definition(db),
@@ -6221,17 +6223,27 @@ impl<'db> Type<'db> {
62216223
.getter(db)
62226224
.and_then(|getter| getter.definition(db))
62236225
.or_else(|| property.setter(db).and_then(|setter| setter.definition(db)))
6224-
.or_else(|| property.deleter(db).and_then(|deleter| deleter.definition(db))),
6226+
.or_else(|| {
6227+
property
6228+
.deleter(db)
6229+
.and_then(|deleter| deleter.definition(db))
6230+
}),
6231+
6232+
Self::LiteralValue(literal) => literal
6233+
.as_enum()
6234+
.and_then(|enum_lit| enum_lit.definition(db))
6235+
.map(TypeDefinition::EnumMember)
6236+
.or_else(|| self.to_meta_type(db).definition(db)),
62256237

6226-
Self::LiteralValue(_)
6227-
// TODO: For enum literals, it would be even better to jump to the definition of the specific member
6228-
| Self::KnownBoundMethod(_)
6238+
Self::KnownBoundMethod(_)
62296239
| Self::WrapperDescriptor(_)
62306240
| Self::DataclassDecorator(_)
62316241
| Self::DataclassTransformer(_)
62326242
| Self::BoundSuper(_) => self.to_meta_type(db).definition(db),
62336243

6234-
Self::TypeVar(bound_typevar) => Some(TypeDefinition::TypeVar(bound_typevar.typevar(db).definition(db)?)),
6244+
Self::TypeVar(bound_typevar) => Some(TypeDefinition::TypeVar(
6245+
bound_typevar.typevar(db).definition(db)?,
6246+
)),
62356247

62366248
Self::ProtocolInstance(protocol) => match protocol.inner {
62376249
Protocol::FromClass(class) => class.type_definition(db),
@@ -6244,8 +6256,12 @@ impl<'db> Type<'db> {
62446256

62456257
Self::SpecialForm(special_form) => special_form.definition(db),
62466258
Self::Never => Type::SpecialForm(SpecialFormType::Never).definition(db),
6247-
Self::Dynamic(DynamicType::Any) => Type::SpecialForm(SpecialFormType::Any).definition(db),
6248-
Self::Dynamic(DynamicType::Unknown | DynamicType::UnknownGeneric(_)) => Type::SpecialForm(SpecialFormType::Unknown).definition(db),
6259+
Self::Dynamic(DynamicType::Any) => {
6260+
Type::SpecialForm(SpecialFormType::Any).definition(db)
6261+
}
6262+
Self::Dynamic(DynamicType::Unknown | DynamicType::UnknownGeneric(_)) => {
6263+
Type::SpecialForm(SpecialFormType::Unknown).definition(db)
6264+
}
62496265
Self::AlwaysTruthy => Type::SpecialForm(SpecialFormType::AlwaysTruthy).definition(db),
62506266
Self::AlwaysFalsy => Type::SpecialForm(SpecialFormType::AlwaysFalsy).definition(db),
62516267

@@ -6257,7 +6273,7 @@ impl<'db> Type<'db> {
62576273
| DynamicType::TodoStarredExpression
62586274
| DynamicType::TodoTypeVarTuple
62596275
| DynamicType::InvalidConcatenateUnknown
6260-
| DynamicType::UnspecializedTypeVar
6276+
| DynamicType::UnspecializedTypeVar,
62616277
)
62626278
| Self::Callable(_)
62636279
| Self::TypeIs(_)

‎crates/ty_python_semantic/src/types/definition.rs‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ pub enum TypeDefinition<'db> {
1818
TypeAlias(Definition<'db>),
1919
NewType(Definition<'db>),
2020
SpecialForm(Definition<'db>),
21+
EnumMember(Definition<'db>),
2122
}
2223

2324
impl TypeDefinition<'_> {
@@ -30,7 +31,8 @@ impl TypeDefinition<'_> {
3031
| Self::TypeVar(definition)
3132
| Self::TypeAlias(definition)
3233
| Self::SpecialForm(definition)
33-
| Self::NewType(definition) => {
34+
| Self::NewType(definition)
35+
| Self::EnumMember(definition) => {
3436
let module = parsed_module(db, definition.file(db)).load(db);
3537
Some(definition.focus_range(db, &module))
3638
}
@@ -50,7 +52,8 @@ impl TypeDefinition<'_> {
5052
| Self::TypeVar(definition)
5153
| Self::TypeAlias(definition)
5254
| Self::SpecialForm(definition)
53-
| Self::NewType(definition) => {
55+
| Self::NewType(definition)
56+
| Self::EnumMember(definition) => {
5457
let module = parsed_module(db, definition.file(db)).load(db);
5558
Some(definition.full_range(db, &module))
5659
}
@@ -66,7 +69,8 @@ impl TypeDefinition<'_> {
6669
| Self::TypeVar(definition)
6770
| Self::TypeAlias(definition)
6871
| Self::SpecialForm(definition)
69-
| Self::NewType(definition) => Some(definition.file(db)),
72+
| Self::NewType(definition)
73+
| Self::EnumMember(definition) => Some(definition.file(db)),
7074
}
7175
}
7276
}
@@ -81,7 +85,8 @@ impl<'db> TypeDefinition<'db> {
8185
| Self::TypeVar(definition)
8286
| Self::TypeAlias(definition)
8387
| Self::SpecialForm(definition)
84-
| Self::NewType(definition) => Some(definition),
88+
| Self::NewType(definition)
89+
| Self::EnumMember(definition) => Some(definition),
8590
}
8691
}
8792
}

‎crates/ty_python_semantic/src/types/display.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1227,7 +1227,12 @@ impl<'db> FmtDetailed<'db> for DisplayRepresentation<'db> {
12271227
.enum_class(self.db)
12281228
.display_with(self.db, self.settings.clone())
12291229
.fmt_detailed(f)?;
1230-
write!(f, ".{}", enum_literal.name(self.db))
1230+
f.write_char('.')?;
1231+
write!(
1232+
f.with_type(Type::enum_literal(enum_literal)),
1233+
"{}",
1234+
enum_literal.name(self.db)
1235+
)
12311236
}
12321237
},
12331238
Type::TypeVar(bound_typevar) => {

‎crates/ty_python_semantic/src/types/literal.rs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ use ruff_python_ast::name::Name;
55
use crate::Db;
66
use crate::types::set_theoretic::RecursivelyDefined;
77
use crate::types::{ClassLiteral, KnownClass, Type};
8+
use ty_python_core::definition::Definition;
9+
use ty_python_core::{place_table, use_def_map};
810

911
/// A literal value. See [`LiteralValueTypeKind`] for details.
1012
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, salsa::Update, get_size2::GetSize)]
@@ -385,4 +387,17 @@ impl<'db> EnumLiteralType<'db> {
385387
pub(crate) fn enum_class_instance(self, db: &'db dyn Db) -> Type<'db> {
386388
self.enum_class(db).to_non_generic_instance(db)
387389
}
390+
391+
pub(crate) fn definition(self, db: &'db dyn Db) -> Option<Definition<'db>> {
392+
let ClassLiteral::Static(class) = self.enum_class(db) else {
393+
return None;
394+
};
395+
396+
let scope = class.body_scope(db);
397+
let symbol_id = place_table(db, scope).symbol_id(self.name(db))?;
398+
399+
use_def_map(db, scope)
400+
.end_of_scope_symbol_bindings(symbol_id)
401+
.find_map(|binding| binding.binding.definition())
402+
}
388403
}

0 commit comments

Comments
 (0)
⚡