Skip to content

Commit b7aad0e

Browse files
rizziprizzip
authored
Merge 613b0da into 72742e7 
2 parents 72742e7 + 613b0da commit b7aad0e

3 files changed

Lines changed: 128 additions & 33 deletions

File tree

‎crates/ty_python_semantic/resources/mdtest/enums.md‎

Lines changed: 48 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -990,9 +990,8 @@ reveal_type(Color.RED._name_) # revealed: Literal["RED"]
990990
def _(red_or_blue: Literal[Color.RED, Color.BLUE]):
991991
reveal_type(red_or_blue.name) # revealed: Literal["RED", "BLUE"]
992992

993-
def _(any_color: Color):
994-
# TODO: Literal["RED", "GREEN", "BLUE"]
995-
reveal_type(any_color.name) # revealed: Any
993+
def _(color: Color):
994+
reveal_type(color.name) # revealed: Literal["RED", "GREEN", "BLUE"]
996995
```
997996

998997
### `value` and `_value_`
@@ -1017,6 +1016,9 @@ reveal_type(Color.RED._value_) # revealed: Literal[1]
10171016
reveal_type(Color.GREEN.value) # revealed: Literal[2]
10181017
reveal_type(Color.GREEN._value_) # revealed: Literal[2]
10191018

1019+
def _(color: Color):
1020+
reveal_type(color.value) # revealed: Literal[1, 2, 3]
1021+
10201022
class Answer(StrEnum):
10211023
YES = "yes"
10221024
NO = "no"
@@ -1026,6 +1028,9 @@ reveal_type(Answer.YES._value_) # revealed: Literal["yes"]
10261028

10271029
reveal_type(Answer.NO.value) # revealed: Literal["no"]
10281030
reveal_type(Answer.NO._value_) # revealed: Literal["no"]
1031+
1032+
def _(answer: Answer):
1033+
reveal_type(answer.value) # revealed: Literal["yes", "no"]
10291034
```
10301035

10311036
## Properties of enum types
@@ -1140,6 +1145,9 @@ python-version = "3.9"
11401145
from enum import Enum, EnumMeta
11411146

11421147
class EnumWithEnumMetaMetaclass(metaclass=EnumMeta):
1148+
# Using `EnumMeta` as a metaclass without inheriting `Enum` requires an `__init__`
1149+
# method that will accept member values (TODO we could catch the lack of this):
1150+
def __init__(self, val): ...
11431151
NO = 0
11441152
YES = 1
11451153

@@ -1148,24 +1156,37 @@ reveal_type(EnumWithEnumMetaMetaclass.NO) # revealed: Literal[EnumWithEnumMetaM
11481156
class SubclassOfEnumMeta(EnumMeta): ...
11491157

11501158
class EnumWithSubclassOfEnumMetaMetaclass(metaclass=SubclassOfEnumMeta):
1159+
def __init__(self, val): ...
11511160
NO = 0
11521161
YES = 1
11531162

11541163
reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO) # revealed: Literal[EnumWithSubclassOfEnumMetaMetaclass.NO]
11551164

1156-
# Attributes like `.value` can *not* be accessed on members of these enums:
1165+
# Attributes `.value` and `.name` can *not* be accessed on members of these enums:
1166+
11571167
# error: [unresolved-attribute]
11581168
EnumWithSubclassOfEnumMetaMetaclass.NO.value
11591169
# error: [unresolved-attribute]
1160-
EnumWithSubclassOfEnumMetaMetaclass.NO._value_
1161-
# error: [unresolved-attribute]
11621170
EnumWithSubclassOfEnumMetaMetaclass.NO.name
1163-
# error: [unresolved-attribute]
1164-
EnumWithSubclassOfEnumMetaMetaclass.NO._name_
1171+
1172+
# But the internal underscore attributes are available:
1173+
1174+
reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO._value_) # revealed: Any
1175+
reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO._name_) # revealed: Literal["NO"]
1176+
1177+
def _(x: EnumWithSubclassOfEnumMetaMetaclass):
1178+
# error: [unresolved-attribute]
1179+
x.value
1180+
# error: [unresolved-attribute]
1181+
x.name
1182+
reveal_type(x._value_) # revealed: Any
1183+
reveal_type(x._name_) # revealed: Literal["NO", "YES"]
11651184
```
11661185

11671186
### Enums with (subclasses of) `EnumType` as metaclass
11681187

1188+
In Python 3.11, the meta-type was renamed to `EnumType`.
1189+
11691190
```toml
11701191
[environment]
11711192
python-version = "3.11"
@@ -1175,6 +1196,7 @@ python-version = "3.11"
11751196
from enum import Enum, EnumType
11761197

11771198
class EnumWithEnumMetaMetaclass(metaclass=EnumType):
1199+
def __init__(self, val): ...
11781200
NO = 0
11791201
YES = 1
11801202

@@ -1183,13 +1205,31 @@ reveal_type(EnumWithEnumMetaMetaclass.NO) # revealed: Literal[EnumWithEnumMetaM
11831205
class SubclassOfEnumMeta(EnumType): ...
11841206

11851207
class EnumWithSubclassOfEnumMetaMetaclass(metaclass=SubclassOfEnumMeta):
1208+
def __init__(self, val): ...
11861209
NO = 0
11871210
YES = 1
11881211

11891212
reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO) # revealed: Literal[EnumWithSubclassOfEnumMetaMetaclass.NO]
11901213

1214+
# Attributes `.value` and `.name` can *not* be accessed on members of these enums:
1215+
11911216
# error: [unresolved-attribute]
11921217
EnumWithSubclassOfEnumMetaMetaclass.NO.value
1218+
# error: [unresolved-attribute]
1219+
EnumWithSubclassOfEnumMetaMetaclass.NO.name
1220+
1221+
# But the internal underscore attributes are available:
1222+
1223+
reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO._value_) # revealed: Any
1224+
reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO._name_) # revealed: Literal["NO"]
1225+
1226+
def _(x: EnumWithSubclassOfEnumMetaMetaclass):
1227+
# error: [unresolved-attribute]
1228+
x.value
1229+
# error: [unresolved-attribute]
1230+
x.name
1231+
reveal_type(x._value_) # revealed: Any
1232+
reveal_type(x._name_) # revealed: Literal["NO", "YES"]
11931233
```
11941234

11951235
## Function syntax

‎crates/ty_python_semantic/src/types.rs‎

Lines changed: 28 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ use crate::types::constraints::{
5858
use crate::types::context::{LintDiagnosticGuard, LintDiagnosticGuardBuilder};
5959
use crate::types::diagnostic::{INVALID_AWAIT, INVALID_TYPE_FORM};
6060
pub use crate::types::display::{DisplaySettings, TypeDetail, TypeDisplayDetails};
61-
use crate::types::enums::{enum_metadata, is_single_member_enum};
61+
use crate::types::enums::enum_metadata;
6262
use crate::types::function::{
6363
DataclassTransformerFlags, DataclassTransformerParams, FunctionDecorators, FunctionSpans,
6464
FunctionType, KnownFunction,
@@ -3308,26 +3308,22 @@ impl<'db> Type<'db> {
33083308
.member_lookup_with_policy(db, name, policy),
33093309

33103310
Type::LiteralValue(literal)
3311-
if literal.as_enum().is_some_and(|enum_literal| {
3312-
matches!(name_str, "name" | "_name_")
3313-
&& Type::ClassLiteral(enum_literal.enum_class(db))
3314-
.is_subtype_of(db, KnownClass::Enum.to_subclass_of(db))
3315-
}) =>
3311+
if literal.as_enum().is_some()
3312+
&& matches!(name_str, "name" | "_name_" | "value" | "_value_") =>
33163313
{
33173314
let enum_literal = literal.as_enum().unwrap();
3318-
Place::bound(Type::string_literal(db, enum_literal.name(db))).into()
3319-
}
3320-
3321-
Type::LiteralValue(literal)
3322-
if literal.as_enum().is_some_and(|enum_literal| {
3323-
matches!(name_str, "value" | "_value_")
3324-
&& Type::ClassLiteral(enum_literal.enum_class(db))
3325-
.is_subtype_of(db, KnownClass::Enum.to_subclass_of(db))
3326-
}) =>
3327-
{
3328-
let enum_literal = literal.as_enum().unwrap();
3329-
enum_metadata(db, enum_literal.enum_class(db))
3330-
.and_then(|metadata| metadata.value_type(enum_literal.name(db)))
3315+
let enum_class = enum_literal.enum_class(db);
3316+
let is_enum_subclass = Type::ClassLiteral(enum_class)
3317+
.is_subtype_of(db, KnownClass::Enum.to_subclass_of(db));
3318+
3319+
enum_metadata(db, enum_class)
3320+
.and_then(|metadata| match name_str {
3321+
"name" if is_enum_subclass => metadata.name_type(db, enum_literal.name(db)),
3322+
"_name_" => metadata.name_type(db, enum_literal.name(db)),
3323+
"value" if is_enum_subclass => metadata.value_type(enum_literal.name(db)),
3324+
"_value_" => metadata.value_type(enum_literal.name(db)),
3325+
_ => None,
3326+
})
33313327
.map_or_else(|| Place::Undefined, Place::bound)
33323328
.into()
33333329
}
@@ -3346,13 +3342,20 @@ impl<'db> Type<'db> {
33463342
}
33473343

33483344
Type::NominalInstance(instance)
3349-
if matches!(name_str, "value" | "_value_")
3350-
&& is_single_member_enum(db, instance.class_literal(db)) =>
3345+
if matches!(name_str, "name" | "_name_" | "value" | "_value_")
3346+
&& enum_metadata(db, instance.class_literal(db)).is_some() =>
33513347
{
3352-
enum_metadata(db, instance.class_literal(db))
3353-
.and_then(|metadata| {
3354-
let (name, _) = metadata.members.get_index(0)?;
3355-
metadata.value_type(name)
3348+
let class_literal = instance.class_literal(db);
3349+
let is_enum_subclass = Type::ClassLiteral(class_literal)
3350+
.is_subtype_of(db, KnownClass::Enum.to_subclass_of(db));
3351+
3352+
enum_metadata(db, class_literal)
3353+
.and_then(|metadata| match name_str {
3354+
"name" if is_enum_subclass => metadata.instance_name_type(db),
3355+
"_name_" => metadata.instance_name_type(db),
3356+
"value" if is_enum_subclass => metadata.instance_value_type(db),
3357+
"_value_" => metadata.instance_value_type(db),
3358+
_ => None,
33563359
})
33573360
.map_or_else(Place::default, Place::bound)
33583361
.into()

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use crate::{
1111
types::{
1212
ClassBase, ClassLiteral, DynamicType, EnumLiteralType, KnownClass, LiteralValueTypeKind,
1313
MemberLookupPolicy, StaticClassLiteral, Type, TypeQualifiers, function::FunctionType,
14+
set_theoretic::builder::UnionBuilder,
1415
},
1516
};
1617

@@ -62,6 +63,57 @@ impl<'db> EnumMetadata<'db> {
6263
}
6364
}
6465

66+
/// Returns the type of `.name`/`._name_` for a given enum member.
67+
///
68+
/// This is always a string literal of the member name.
69+
pub(crate) fn name_type(&self, db: &'db dyn Db, member_name: &Name) -> Option<Type<'db>> {
70+
self.members
71+
.contains_key(member_name)
72+
.then(|| Type::string_literal(db, member_name.as_str()))
73+
}
74+
75+
/// Returns the type of `.value`/`._value_` for an enum instance that is not
76+
/// narrowed to a specific member (e.g. `x: MyEnum` where `MyEnum` has multiple members).
77+
///
78+
/// If there is an explicit `_value_` annotation, returns that.
79+
/// If there is a custom `__init__`, returns `Any`.
80+
/// Otherwise, returns the union of all member value types.
81+
pub(crate) fn instance_value_type(&self, db: &'db dyn Db) -> Option<Type<'db>> {
82+
if self.members.is_empty() {
83+
return None;
84+
}
85+
if let Some(annotation) = self.value_annotation {
86+
Some(annotation)
87+
} else if self.init_function.is_some() {
88+
Some(Type::Dynamic(DynamicType::Any))
89+
} else {
90+
let union = self
91+
.members
92+
.values()
93+
.copied()
94+
.fold(UnionBuilder::new(db), UnionBuilder::add)
95+
.build();
96+
Some(union)
97+
}
98+
}
99+
100+
/// Returns the type of `.name`/`._name_` for an enum instance that is not
101+
/// narrowed to a specific member (e.g. `x: MyEnum` where `MyEnum` has multiple members).
102+
///
103+
/// Returns the union of all member name string literals.
104+
pub(crate) fn instance_name_type(&self, db: &'db dyn Db) -> Option<Type<'db>> {
105+
if self.members.is_empty() {
106+
return None;
107+
}
108+
let union = self
109+
.members
110+
.keys()
111+
.map(|name| Type::string_literal(db, name.as_str()))
112+
.fold(UnionBuilder::new(db), UnionBuilder::add)
113+
.build();
114+
Some(union)
115+
}
116+
65117
pub(crate) fn resolve_member<'a>(&'a self, name: &'a Name) -> Option<&'a Name> {
66118
if self.members.contains_key(name) {
67119
Some(name)

0 commit comments

Comments
 (0)