Skip to content

Commit a90e404

Browse files
authored
[red-knot] PEP 695 type aliases (#14357)
## Summary Add support for (non-generic) type aliases. The main motivation behind this was to get rid of panics involving expressions in (generic) type aliases. But it turned out the best way to fix it was to implement (partial) support for type aliases. ```py type IntOrStr = int | str reveal_type(IntOrStr) # revealed: typing.TypeAliasType reveal_type(IntOrStr.__name__) # revealed: Literal["IntOrStr"] x: IntOrStr = 1 reveal_type(x) # revealed: Literal[1] def f() -> None: reveal_type(x) # revealed: int | str ``` ## Test Plan - Updated corpus test allow list to reflect that we don't panic anymore. - Added Markdown-based test for type aliases (`type_alias.md`)
1 parent 8358ad8 commit a90e404

File tree

9 files changed

+298
-51
lines changed

9 files changed

+298
-51
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
# Type aliases
2+
3+
## Basic
4+
5+
```py
6+
type IntOrStr = int | str
7+
8+
reveal_type(IntOrStr) # revealed: typing.TypeAliasType
9+
reveal_type(IntOrStr.__name__) # revealed: Literal["IntOrStr"]
10+
11+
x: IntOrStr = 1
12+
13+
reveal_type(x) # revealed: Literal[1]
14+
15+
def f() -> None:
16+
reveal_type(x) # revealed: int | str
17+
```
18+
19+
## `__value__` attribute
20+
21+
```py
22+
type IntOrStr = int | str
23+
24+
# TODO: This should either fall back to the specified type from typeshed,
25+
# which is `Any`, or be the actual type of the runtime value expression
26+
# `int | str`, i.e. `types.UnionType`.
27+
reveal_type(IntOrStr.__value__) # revealed: @Todo(instance attributes)
28+
```
29+
30+
## Invalid assignment
31+
32+
```py
33+
type OptionalInt = int | None
34+
35+
# error: [invalid-assignment]
36+
x: OptionalInt = "1"
37+
```
38+
39+
## Type aliases in type aliases
40+
41+
```py
42+
type IntOrStr = int | str
43+
type IntOrStrOrBytes = IntOrStr | bytes
44+
45+
x: IntOrStrOrBytes = 1
46+
47+
def f() -> None:
48+
reveal_type(x) # revealed: int | str | bytes
49+
```
50+
51+
## Aliased type aliases
52+
53+
```py
54+
type IntOrStr = int | str
55+
MyIntOrStr = IntOrStr
56+
57+
x: MyIntOrStr = 1
58+
59+
# error: [invalid-assignment]
60+
y: MyIntOrStr = None
61+
```
62+
63+
## Generic type aliases
64+
65+
```py
66+
type ListOrSet[T] = list[T] | set[T]
67+
68+
# TODO: Should be `tuple[typing.TypeVar | typing.ParamSpec | typing.TypeVarTuple, ...]`,
69+
# as specified in the `typeshed` stubs.
70+
reveal_type(ListOrSet.__type_params__) # revealed: @Todo(instance attributes)
71+
```

crates/red_knot_python_semantic/src/semantic_index/builder.rs

+23-1
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,8 @@ impl<'db> SemanticIndexBuilder<'db> {
131131
let scope_id = ScopeId::new(self.db, self.file, file_scope_id, countme::Count::default());
132132

133133
self.scope_ids_by_scope.push(scope_id);
134-
self.scopes_by_node.insert(node.node_key(), file_scope_id);
134+
let previous = self.scopes_by_node.insert(node.node_key(), file_scope_id);
135+
debug_assert_eq!(previous, None);
135136

136137
debug_assert_eq!(ast_id_scope, file_scope_id);
137138

@@ -589,6 +590,27 @@ where
589590
},
590591
);
591592
}
593+
ast::Stmt::TypeAlias(type_alias) => {
594+
let symbol = self.add_symbol(
595+
type_alias
596+
.name
597+
.as_name_expr()
598+
.map(|name| name.id.clone())
599+
.unwrap_or("<unknown>".into()),
600+
);
601+
self.add_definition(symbol, type_alias);
602+
self.visit_expr(&type_alias.name);
603+
604+
self.with_type_params(
605+
NodeWithScopeRef::TypeAliasTypeParameters(type_alias),
606+
type_alias.type_params.as_ref(),
607+
|builder| {
608+
builder.push_scope(NodeWithScopeRef::TypeAlias(type_alias));
609+
builder.visit_expr(&type_alias.value);
610+
builder.pop_scope()
611+
},
612+
);
613+
}
592614
ast::Stmt::Import(node) => {
593615
for alias in &node.names {
594616
let symbol_name = if let Some(asname) = &alias.asname {

crates/red_knot_python_semantic/src/semantic_index/definition.rs

+19
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ pub(crate) enum DefinitionNodeRef<'a> {
8383
For(ForStmtDefinitionNodeRef<'a>),
8484
Function(&'a ast::StmtFunctionDef),
8585
Class(&'a ast::StmtClassDef),
86+
TypeAlias(&'a ast::StmtTypeAlias),
8687
NamedExpression(&'a ast::ExprNamed),
8788
Assignment(AssignmentDefinitionNodeRef<'a>),
8889
AnnotatedAssignment(&'a ast::StmtAnnAssign),
@@ -109,6 +110,12 @@ impl<'a> From<&'a ast::StmtClassDef> for DefinitionNodeRef<'a> {
109110
}
110111
}
111112

113+
impl<'a> From<&'a ast::StmtTypeAlias> for DefinitionNodeRef<'a> {
114+
fn from(node: &'a ast::StmtTypeAlias) -> Self {
115+
Self::TypeAlias(node)
116+
}
117+
}
118+
112119
impl<'a> From<&'a ast::ExprNamed> for DefinitionNodeRef<'a> {
113120
fn from(node: &'a ast::ExprNamed) -> Self {
114121
Self::NamedExpression(node)
@@ -265,6 +272,9 @@ impl<'db> DefinitionNodeRef<'db> {
265272
DefinitionNodeRef::Class(class) => {
266273
DefinitionKind::Class(AstNodeRef::new(parsed, class))
267274
}
275+
DefinitionNodeRef::TypeAlias(type_alias) => {
276+
DefinitionKind::TypeAlias(AstNodeRef::new(parsed, type_alias))
277+
}
268278
DefinitionNodeRef::NamedExpression(named) => {
269279
DefinitionKind::NamedExpression(AstNodeRef::new(parsed, named))
270280
}
@@ -358,6 +368,7 @@ impl<'db> DefinitionNodeRef<'db> {
358368
}
359369
Self::Function(node) => node.into(),
360370
Self::Class(node) => node.into(),
371+
Self::TypeAlias(node) => node.into(),
361372
Self::NamedExpression(node) => node.into(),
362373
Self::Assignment(AssignmentDefinitionNodeRef {
363374
value: _,
@@ -434,6 +445,7 @@ pub enum DefinitionKind<'db> {
434445
ImportFrom(ImportFromDefinitionKind),
435446
Function(AstNodeRef<ast::StmtFunctionDef>),
436447
Class(AstNodeRef<ast::StmtClassDef>),
448+
TypeAlias(AstNodeRef<ast::StmtTypeAlias>),
437449
NamedExpression(AstNodeRef<ast::ExprNamed>),
438450
Assignment(AssignmentDefinitionKind<'db>),
439451
AnnotatedAssignment(AstNodeRef<ast::StmtAnnAssign>),
@@ -456,6 +468,7 @@ impl DefinitionKind<'_> {
456468
// functions, classes, and imports always bind, and we consider them declarations
457469
DefinitionKind::Function(_)
458470
| DefinitionKind::Class(_)
471+
| DefinitionKind::TypeAlias(_)
459472
| DefinitionKind::Import(_)
460473
| DefinitionKind::ImportFrom(_)
461474
| DefinitionKind::TypeVar(_)
@@ -682,6 +695,12 @@ impl From<&ast::StmtClassDef> for DefinitionNodeKey {
682695
}
683696
}
684697

698+
impl From<&ast::StmtTypeAlias> for DefinitionNodeKey {
699+
fn from(node: &ast::StmtTypeAlias) -> Self {
700+
Self(NodeKey::from_node(node))
701+
}
702+
}
703+
685704
impl From<&ast::ExprName> for DefinitionNodeKey {
686705
fn from(node: &ast::ExprName) -> Self {
687706
Self(NodeKey::from_node(node))

crates/red_knot_python_semantic/src/semantic_index/symbol.rs

+42-11
Original file line numberDiff line numberDiff line change
@@ -116,14 +116,11 @@ impl<'db> ScopeId<'db> {
116116
// Type parameter scopes behave like function scopes in terms of name resolution; CPython
117117
// symbol table also uses the term "function-like" for these scopes.
118118
matches!(
119-
self.node(db),
120-
NodeWithScopeKind::ClassTypeParameters(_)
121-
| NodeWithScopeKind::FunctionTypeParameters(_)
122-
| NodeWithScopeKind::Function(_)
123-
| NodeWithScopeKind::ListComprehension(_)
124-
| NodeWithScopeKind::SetComprehension(_)
125-
| NodeWithScopeKind::DictComprehension(_)
126-
| NodeWithScopeKind::GeneratorExpression(_)
119+
self.node(db).scope_kind(),
120+
ScopeKind::Annotation
121+
| ScopeKind::Function
122+
| ScopeKind::TypeAlias
123+
| ScopeKind::Comprehension
127124
)
128125
}
129126

@@ -144,6 +141,12 @@ impl<'db> ScopeId<'db> {
144141
}
145142
NodeWithScopeKind::Function(function)
146143
| NodeWithScopeKind::FunctionTypeParameters(function) => function.name.as_str(),
144+
NodeWithScopeKind::TypeAlias(type_alias)
145+
| NodeWithScopeKind::TypeAliasTypeParameters(type_alias) => type_alias
146+
.name
147+
.as_name_expr()
148+
.map(|name| name.id.as_str())
149+
.unwrap_or("<type alias>"),
147150
NodeWithScopeKind::Lambda(_) => "<lambda>",
148151
NodeWithScopeKind::ListComprehension(_) => "<listcomp>",
149152
NodeWithScopeKind::SetComprehension(_) => "<setcomp>",
@@ -201,6 +204,7 @@ pub enum ScopeKind {
201204
Class,
202205
Function,
203206
Comprehension,
207+
TypeAlias,
204208
}
205209

206210
impl ScopeKind {
@@ -326,6 +330,8 @@ pub(crate) enum NodeWithScopeRef<'a> {
326330
Lambda(&'a ast::ExprLambda),
327331
FunctionTypeParameters(&'a ast::StmtFunctionDef),
328332
ClassTypeParameters(&'a ast::StmtClassDef),
333+
TypeAlias(&'a ast::StmtTypeAlias),
334+
TypeAliasTypeParameters(&'a ast::StmtTypeAlias),
329335
ListComprehension(&'a ast::ExprListComp),
330336
SetComprehension(&'a ast::ExprSetComp),
331337
DictComprehension(&'a ast::ExprDictComp),
@@ -347,6 +353,12 @@ impl NodeWithScopeRef<'_> {
347353
NodeWithScopeRef::Function(function) => {
348354
NodeWithScopeKind::Function(AstNodeRef::new(module, function))
349355
}
356+
NodeWithScopeRef::TypeAlias(type_alias) => {
357+
NodeWithScopeKind::TypeAlias(AstNodeRef::new(module, type_alias))
358+
}
359+
NodeWithScopeRef::TypeAliasTypeParameters(type_alias) => {
360+
NodeWithScopeKind::TypeAliasTypeParameters(AstNodeRef::new(module, type_alias))
361+
}
350362
NodeWithScopeRef::Lambda(lambda) => {
351363
NodeWithScopeKind::Lambda(AstNodeRef::new(module, lambda))
352364
}
@@ -387,6 +399,12 @@ impl NodeWithScopeRef<'_> {
387399
NodeWithScopeRef::ClassTypeParameters(class) => {
388400
NodeWithScopeKey::ClassTypeParameters(NodeKey::from_node(class))
389401
}
402+
NodeWithScopeRef::TypeAlias(type_alias) => {
403+
NodeWithScopeKey::TypeAlias(NodeKey::from_node(type_alias))
404+
}
405+
NodeWithScopeRef::TypeAliasTypeParameters(type_alias) => {
406+
NodeWithScopeKey::TypeAliasTypeParameters(NodeKey::from_node(type_alias))
407+
}
390408
NodeWithScopeRef::ListComprehension(comprehension) => {
391409
NodeWithScopeKey::ListComprehension(NodeKey::from_node(comprehension))
392410
}
@@ -411,6 +429,8 @@ pub enum NodeWithScopeKind {
411429
ClassTypeParameters(AstNodeRef<ast::StmtClassDef>),
412430
Function(AstNodeRef<ast::StmtFunctionDef>),
413431
FunctionTypeParameters(AstNodeRef<ast::StmtFunctionDef>),
432+
TypeAliasTypeParameters(AstNodeRef<ast::StmtTypeAlias>),
433+
TypeAlias(AstNodeRef<ast::StmtTypeAlias>),
414434
Lambda(AstNodeRef<ast::ExprLambda>),
415435
ListComprehension(AstNodeRef<ast::ExprListComp>),
416436
SetComprehension(AstNodeRef<ast::ExprSetComp>),
@@ -423,9 +443,11 @@ impl NodeWithScopeKind {
423443
match self {
424444
Self::Module => ScopeKind::Module,
425445
Self::Class(_) => ScopeKind::Class,
426-
Self::Function(_) => ScopeKind::Function,
427-
Self::Lambda(_) => ScopeKind::Function,
428-
Self::FunctionTypeParameters(_) | Self::ClassTypeParameters(_) => ScopeKind::Annotation,
446+
Self::Function(_) | Self::Lambda(_) => ScopeKind::Function,
447+
Self::FunctionTypeParameters(_)
448+
| Self::ClassTypeParameters(_)
449+
| Self::TypeAliasTypeParameters(_) => ScopeKind::Annotation,
450+
Self::TypeAlias(_) => ScopeKind::TypeAlias,
429451
Self::ListComprehension(_)
430452
| Self::SetComprehension(_)
431453
| Self::DictComprehension(_)
@@ -446,6 +468,13 @@ impl NodeWithScopeKind {
446468
_ => panic!("expected function"),
447469
}
448470
}
471+
472+
pub fn expect_type_alias(&self) -> &ast::StmtTypeAlias {
473+
match self {
474+
Self::TypeAlias(type_alias) => type_alias.node(),
475+
_ => panic!("expected type alias"),
476+
}
477+
}
449478
}
450479

451480
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
@@ -455,6 +484,8 @@ pub(crate) enum NodeWithScopeKey {
455484
ClassTypeParameters(NodeKey),
456485
Function(NodeKey),
457486
FunctionTypeParameters(NodeKey),
487+
TypeAlias(NodeKey),
488+
TypeAliasTypeParameters(NodeKey),
458489
Lambda(NodeKey),
459490
ListComprehension(NodeKey),
460491
SetComprehension(NodeKey),

0 commit comments

Comments
 (0)