Skip to content

Commit

Permalink
Handle inferring Literal type annotation
Browse files Browse the repository at this point in the history
  • Loading branch information
Glyphack committed Oct 27, 2024
1 parent 855405c commit a5a4f7f
Show file tree
Hide file tree
Showing 5 changed files with 111 additions and 57 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -95,13 +95,11 @@ class AlwaysFalse:
def __contains__(self, item: int) -> Literal[""]:
return ""

# TODO: it should be Literal[True] and Literal[False]
reveal_type(42 in AlwaysTrue()) # revealed: @Todo
reveal_type(42 not in AlwaysTrue()) # revealed: @Todo
reveal_type(42 in AlwaysTrue()) # revealed: Literal[True]
reveal_type(42 not in AlwaysTrue()) # revealed: Literal[False]

# TODO: it should be Literal[False] and Literal[True]
reveal_type(42 in AlwaysFalse()) # revealed: @Todo
reveal_type(42 not in AlwaysFalse()) # revealed: @Todo
reveal_type(42 in AlwaysFalse()) # revealed: Literal[False]
reveal_type(42 not in AlwaysFalse()) # revealed: Literal[True]
```

## No Fallback for `__contains__`
Expand Down
43 changes: 28 additions & 15 deletions crates/red_knot_python_semantic/resources/mdtest/literal/literal.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,24 +3,37 @@
```py
from typing import Literal

x: Literal[10] = 10
mode: Literal["w", "r"] = "w"
mode2: Literal["w"] | Literal["r"] = "w"

# reveal_type(Literal[26]) # revealed: Literal[26]
# reveal_type(Literal[0x1A]) # revealed: Literal[26]
# reveal_type(Literal[-4]) # revealed: Literal[-4]
# reveal_type(Literal["hello world"]) # revealed: Literal["hello world"]
# reveal_type(Literal[b"hello world"]) # revealed: Literal[b"hello world"]
# reveal_type(Literal["hello world"]) # revealed: Literal["hello world"]
# reveal_type(Literal[True]) # revealed: Literal[True]
# reveal_type(Literal[Color.RED]) # revealed: Literal["Red"]
# reveal_type(Literal[None]) # revealed: None
# TODO: "Revealed type is `Literal[1, 2, 3] | Literal["foo"] | Literal[5] | None`"
# reveal_type(Literal[Literal[Literal[1, 2, 3], "foo"], 5, None]) # revealed: Literal[1, 2, 3, "foo", 5, None]
# TODO: PEP-604 unions should not give error
# error: [unsupported-operator] "Operator `|` is unsupported between objects of type `object` and `object`"
mode2: Literal["w"] | Literal["r"] = "w"
# union_var: Literal[Literal[Literal[1, 2, 3], "foo"], 5, None]
# reveal_type(union_var) # revealed: Literal[1, 2, 3, "foo", 5, None]

def f():
reveal_type(x) # revealed: Literal[10]
reveal_type(mode) # revealed: Literal["w", "r"]
# reveal_type(mode2) # revealed: Literal["w", "r"]
reveal_type(mode2) # revealed: @Todo

a: Literal[26] = 26
reveal_type(a) # revealed: Literal[26]
a2: Literal[0x1A] = 0x1A
reveal_type(a2) # revealed: Literal[26]
a3: Literal[-4] = -4
reveal_type(a3) # revealed: Literal[-4]
a4: Literal["hello world"] = "hello world"
reveal_type(a4) # revealed: Literal["hello world"]
a5: Literal[b"hello world"] = b"hello world"
reveal_type(a5) # revealed: Literal[b"hello world"]
a6: Literal["hello world"] = "hello world"
reveal_type(a6) # revealed: Literal["hello world"]
a7: Literal[True] = True
reveal_type(a7) # revealed: Literal[True]
# a: Literal[Color.RED]
# reveal_type(a)
a8: Literal[None] = None
reveal_type(a8) # revealed: None

# error: [invalid-literal-parameter] "Type arguments for `Literal` must be None, a literal value (int, bool, str, or bytes), or an enum value"
a9: Literal[3 + 4] = 7
```
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Unary Operations

```py
from typing import Literal

class Number:
def __init__(self, value: int):
self.value = 1
Expand All @@ -19,7 +21,7 @@ a = Number()
reveal_type(+a) # revealed: int
reveal_type(-a) # revealed: int
# TODO: this should be True
reveal_type(~a) # revealed: object
reveal_type(~a) # revealed: Literal[True]

class NoDunder: ...

Expand Down
16 changes: 15 additions & 1 deletion crates/red_knot_python_semantic/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1047,10 +1047,16 @@ impl<'db> Type<'db> {
Type::Intersection(_) => Type::Todo,
// TODO: calling `.to_instance()` on any of these should result in a diagnostic,
// since they already indicate that the object is an instance of some kind:
Type::Instance(instance) => {
if instance.known_instance(db).is_some() {
*self
} else {
Type::Unknown
}
}
Type::BooleanLiteral(_)
| Type::BytesLiteral(_)
| Type::FunctionLiteral(_)
| Type::Instance(_)
| Type::ModuleLiteral(_)
| Type::IntLiteral(_)
| Type::StringLiteral(_)
Expand Down Expand Up @@ -1279,6 +1285,14 @@ pub enum KnownInstance {
Literal,
}

impl KnownInstance {
pub const fn as_str(&self) -> &'static str {
match self {
KnownInstance::Literal => "Literal",
}
}
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum CallOutcome<'db> {
Callable {
Expand Down
95 changes: 61 additions & 34 deletions crates/red_knot_python_semantic/src/types/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1380,19 +1380,13 @@ impl<'db> TypeInferenceBuilder<'db> {
let class = instance.class_type(self.db);
if class.is_known(self.db, KnownClass::SpecialForm) {
if let Some(name_expr) = target.as_name_expr() {
let target_name = name_expr.id.clone();
let new_class = ClassType::new(
self.db,
target_name,
class.definition(self.db),
class.body_scope(self.db),
Some(KnownClass::SpecialForm),
);
annotation_ty = Type::Instance(InstanceType::new(
self.db,
new_class,
Some(KnownInstance::Literal),
));
if name_expr.id == KnownInstance::Literal.as_str() {
annotation_ty = Type::Instance(InstanceType::new(
self.db,
class,
Some(KnownInstance::Literal),
));
}
}
}
}
Expand Down Expand Up @@ -3477,29 +3471,60 @@ impl<'db> TypeInferenceBuilder<'db> {
} = subscript;

let value_ty = self.infer_type_expression(value);
let slice_ty = self.infer_type_expression(slice);

// Handling of Special Forms
match (value_ty, slice_ty) {
(Type::Instance(class), slice_ty)
if class.is_known(self.db, KnownClass::SpecialForm) =>
{
match class.known_instance(self.db) {
Some(s) => match s {
KnownInstance::Literal => match slice_ty {
// TODO: not correct logic
match value_ty {
Type::Instance(class) if class.is_known(self.db, KnownClass::SpecialForm) => {
let Some(s) = class.known_instance(self.db) else {
return Type::Todo;
};
// NOTE: slice_ty is treated as expression because Literal accepts expression
// inside the []
let slice_ty = self.infer_expression(slice);
match s {
KnownInstance::Literal => {
match **slice {
ruff_python_ast::Expr::StringLiteral(_)
| ruff_python_ast::Expr::BytesLiteral(_)
| ruff_python_ast::Expr::NumberLiteral(_)
| ruff_python_ast::Expr::BooleanLiteral(_)
| ruff_python_ast::Expr::Tuple(_)
// For enum values
| ruff_python_ast::Expr::Attribute(_)
// For Another Literal inside this Literal
| ruff_python_ast::Expr::Subscript(_)
| ruff_python_ast::Expr::NoneLiteral(_) => {}
// for negative numbers
ruff_python_ast::Expr::UnaryOp(ref u) if (u.op == UnaryOp::USub || u.op == UnaryOp::UAdd) && u.operand.is_number_literal_expr() => {}
_ => {
self.add_diagnostic(
slice.as_ref().into(),
"invalid-literal-parameter",
format_args!(
"Type arguments for `Literal` must be None, a literal value (int, bool, str, or bytes), or an enum value",
),
);
return Type::Unknown;
}
};
match slice_ty {
Type::Tuple(tuple) => {
let elts = tuple.elements(self.db);
Type::Union(UnionType::new(self.db, elts))
}
ty => ty,
},
},
None => todo!(),
}
}
}
}
// TODO: emit diagnostic
_ => Type::Todo,
value_ty => {
let value_node = value.as_ref();
let slice_ty = self.infer_expression(slice);
// TODO: currently the logic to get the type of type of a subscript in type
// annotation with where value is a slice is defined in the
// subscript_expression_types. Once it's complete call that logic from here.
self.infer_subscript_expression_types(value_node, value_ty, slice_ty)
}
}
}
}
Expand Down Expand Up @@ -3528,8 +3553,6 @@ impl<'db> TypeInferenceBuilder<'db> {
impl<'db> TypeInferenceBuilder<'db> {
fn infer_type_expression(&mut self, expression: &ast::Expr) -> Type<'db> {
// https://typing.readthedocs.io/en/latest/spec/annotations.html#grammar-token-expression-grammar-type_expression
// TODO: this does not include any of the special forms, and is only a
// stub of the forms other than a standalone name in scope.

let ty = match expression {
ast::Expr::Name(name) => {
Expand Down Expand Up @@ -3835,13 +3858,17 @@ fn perform_rich_comparison<'db>(
// TODO: this currently gives the return type even if the arg types are invalid
// (e.g. int.__lt__ with string instance should be errored, currently bool)

let left_instance = Type::Instance(InstanceType::new(db, left_class, None));
let right_instance = Type::Instance(InstanceType::new(db, right_class, None));
let call_dunder =
|op: RichCompareOperator, left_class: ClassType<'db>, right_class: ClassType<'db>| {
left_class
.class_member(db, op.dunder())
.call(db, &[left_instance, right_instance])
.call(
db,
&[
Type::Instance(InstanceType::new(db, left_class, None)),
Type::Instance(InstanceType::new(db, right_class, None)),
],
)
.return_ty(db)
};

Expand All @@ -3865,8 +3892,8 @@ fn perform_rich_comparison<'db>(
})
.ok_or_else(|| CompareUnsupportedError {
op: op.into(),
left_ty: left_instance,
right_ty: right_instance,
left_ty: Type::Instance(InstanceType::new(db, left_class, None)),
right_ty: Type::Instance(InstanceType::new(db, right_class, None)),
})
}

Expand Down

0 comments on commit a5a4f7f

Please sign in to comment.