Skip to content

Commit

Permalink
Added bytes type and some inference (#13061)
Browse files Browse the repository at this point in the history
## Summary

This PR adds the `bytes` type to red-knot:
- Added the `bytes` type
- Added support for bytes literals
- Support for the `+` operator

Improves on #12701 

Big TODO on supporting and normalizing r-prefixed bytestrings
(`rb"hello\n"`)

## Test Plan

Added a test for a bytes literals, concatenation, and corner values
  • Loading branch information
teofr authored Aug 22, 2024
1 parent 2edd32a commit b9c8113
Show file tree
Hide file tree
Showing 6 changed files with 73 additions and 6 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/red_knot_python_semantic/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ ruff_python_ast = { workspace = true }
ruff_python_stdlib = { workspace = true }
ruff_source_file = { workspace = true }
ruff_text_size = { workspace = true }
ruff_python_literal = { workspace = true }

anyhow = { workspace = true }
bitflags = { workspace = true }
Expand Down
12 changes: 12 additions & 0 deletions crates/red_knot_python_semantic/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,8 @@ pub enum Type<'db> {
IntLiteral(i64),
/// A boolean literal, either `True` or `False`.
BooleanLiteral(bool),
/// A bytes literal
BytesLiteral(BytesLiteralType<'db>),
// TODO protocols, callable types, overloads, generics, type vars
}

Expand Down Expand Up @@ -276,6 +278,10 @@ impl<'db> Type<'db> {
Type::Unknown
}
Type::BooleanLiteral(_) => Type::Unknown,
Type::BytesLiteral(_) => {
// TODO defer to Type::Instance(<bytes from typeshed>).member
Type::Unknown
}
}
}

Expand Down Expand Up @@ -372,6 +378,12 @@ pub struct IntersectionType<'db> {
negative: FxOrderSet<Type<'db>>,
}

#[salsa::interned]
pub struct BytesLiteralType<'db> {
#[return_ref]
value: Box<[u8]>,
}

#[cfg(test)]
mod tests {
use anyhow::Context;
Expand Down
11 changes: 11 additions & 0 deletions crates/red_knot_python_semantic/src/types/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@
use std::fmt::{Display, Formatter};

use ruff_python_ast::str::Quote;
use ruff_python_literal::escape::AsciiEscape;

use crate::types::{IntersectionType, Type, UnionType};
use crate::Db;

Expand Down Expand Up @@ -38,6 +41,14 @@ impl Display for DisplayType<'_> {
Type::BooleanLiteral(boolean) => {
write!(f, "Literal[{}]", if *boolean { "True" } else { "False" })
}
Type::BytesLiteral(bytes) => {
let escape =
AsciiEscape::with_preferred_quote(bytes.value(self.db).as_ref(), Quote::Double);

f.write_str("Literal[")?;
escape.bytes_repr().write(f)?;
f.write_str("]")
}
}
}
}
Expand Down
52 changes: 47 additions & 5 deletions crates/red_knot_python_semantic/src/types/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,8 @@ use crate::semantic_index::symbol::{FileScopeId, NodeWithScopeKind, NodeWithScop
use crate::semantic_index::SemanticIndex;
use crate::types::diagnostic::{TypeCheckDiagnostic, TypeCheckDiagnostics};
use crate::types::{
builtins_symbol_ty_by_name, definitions_ty, global_symbol_ty_by_name, ClassType, FunctionType,
Name, Type, UnionBuilder,
builtins_symbol_ty_by_name, definitions_ty, global_symbol_ty_by_name, BytesLiteralType,
ClassType, FunctionType, Name, Type, UnionBuilder,
};
use crate::Db;

Expand Down Expand Up @@ -1206,9 +1206,12 @@ impl<'db> TypeInferenceBuilder<'db> {
}

#[allow(clippy::unused_self)]
fn infer_bytes_literal_expression(&mut self, _literal: &ast::ExprBytesLiteral) -> Type<'db> {
// TODO
Type::Unknown
fn infer_bytes_literal_expression(&mut self, literal: &ast::ExprBytesLiteral) -> Type<'db> {
// TODO: ignoring r/R prefixes for now, should normalize bytes values
Type::BytesLiteral(BytesLiteralType::new(
self.db,
literal.value.bytes().collect(),
))
}

fn infer_fstring_expression(&mut self, fstring: &ast::ExprFString) -> Type<'db> {
Expand Down Expand Up @@ -1684,6 +1687,7 @@ impl<'db> TypeInferenceBuilder<'db> {
let left_ty = self.infer_expression(left);
let right_ty = self.infer_expression(right);

// TODO flatten the matches by matching on (left_ty, right_ty, op)
match left_ty {
Type::Any => Type::Any,
Type::Unknown => Type::Unknown,
Expand Down Expand Up @@ -1722,6 +1726,22 @@ impl<'db> TypeInferenceBuilder<'db> {
_ => Type::Unknown, // TODO
}
}
Type::BytesLiteral(lhs) => {
match right_ty {
Type::BytesLiteral(rhs) => {
match op {
ast::Operator::Add => Type::BytesLiteral(BytesLiteralType::new(
self.db,
[lhs.value(self.db).as_ref(), rhs.value(self.db).as_ref()]
.concat()
.into_boxed_slice(),
)),
_ => Type::Unknown, // TODO
}
}
_ => Type::Unknown, // TODO
}
}
_ => Type::Unknown, // TODO
}
}
Expand Down Expand Up @@ -2235,6 +2255,28 @@ mod tests {
Ok(())
}

#[test]
fn bytes_type() -> anyhow::Result<()> {
let mut db = setup_db();

db.write_dedented(
"src/a.py",
"
w = b'red' b'knot'
x = b'hello'
y = b'world' + b'!'
z = b'\\xff\\x00'
",
)?;

assert_public_ty(&db, "src/a.py", "w", "Literal[b\"redknot\"]");
assert_public_ty(&db, "src/a.py", "x", "Literal[b\"hello\"]");
assert_public_ty(&db, "src/a.py", "y", "Literal[b\"world!\"]");
assert_public_ty(&db, "src/a.py", "z", "Literal[b\"\\xff\\x00\"]");

Ok(())
}

#[test]
fn resolve_union() -> anyhow::Result<()> {
let mut db = setup_db();
Expand Down
2 changes: 1 addition & 1 deletion crates/ruff_python_ast/src/nodes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2152,7 +2152,7 @@ impl BytesLiteralValue {
}

/// Returns an iterator over the bytes of the concatenated bytes.
fn bytes(&self) -> impl Iterator<Item = u8> + '_ {
pub fn bytes(&self) -> impl Iterator<Item = u8> + '_ {
self.iter().flat_map(|part| part.as_slice().iter().copied())
}
}
Expand Down

0 comments on commit b9c8113

Please sign in to comment.