Skip to main content

ty_python_core/
unpack.rs

1use crate::Program;
2use ruff_db::PythonFile;
3use ruff_db::files::File;
4use ruff_db::parsed::ParsedModuleRef;
5use ruff_python_ast::{self as ast, AnyNodeRef};
6use ruff_text_size::{Ranged, TextRange};
7
8use crate::Db;
9use crate::EvaluationMode;
10use crate::ProgramFile;
11use crate::ast_node_ref::AstNodeRef;
12use crate::expression::Expression;
13use crate::scope::{FileScopeId, ScopeId};
14
15/// This ingredient represents a single unpacking.
16///
17/// This is required to make use of salsa to cache the complete unpacking of multiple variables
18/// involved. It allows us to:
19/// 1. Avoid doing structural match multiple times for each definition
20/// 2. Avoid highlighting the same error multiple times
21///
22/// ## Module-local type
23/// This type should not be used as part of any cross-module API because
24/// it holds a reference to the AST node. Range-offset changes
25/// then propagate through all usages, and deserialization requires
26/// reparsing the entire module.
27///
28/// E.g. don't use this type in:
29///
30/// * a return type of a cross-module query
31/// * a field of a type that is a return type of a cross-module query
32/// * an argument of a cross-module query
33#[salsa::tracked(debug, heap_size=ruff_memory_usage::heap_size)]
34pub struct Unpack<'db> {
35    #[returns(copy)]
36    pub program_file: ProgramFile<'db>,
37
38    #[returns(copy)]
39    pub(crate) value_file_scope: FileScopeId,
40
41    #[returns(copy)]
42    pub(crate) target_file_scope: FileScopeId,
43
44    /// The target expression that is being unpacked. For example, in `(a, b) = (1, 2)`, the target
45    /// expression is `(a, b)`.
46    #[no_eq]
47    #[tracked]
48    #[returns(ref)]
49    pub(crate) _target: AstNodeRef<ast::Expr>,
50
51    /// The ingredient representing the value expression of the unpacking. For example, in
52    /// `(a, b) = (1, 2)`, the value expression is `(1, 2)`.
53    #[returns(copy)]
54    pub value: UnpackValue<'db>,
55}
56
57// The Salsa heap is tracked separately.
58impl get_size2::GetSize for Unpack<'_> {}
59
60impl<'db> Unpack<'db> {
61    pub fn file(self, db: &'db dyn Db) -> File {
62        self.program_file(db).file(db)
63    }
64
65    pub fn python_file(self, db: &'db dyn Db) -> PythonFile<'db> {
66        self.program_file(db).python_file(db)
67    }
68
69    pub fn target<'ast>(self, db: &'db dyn Db, parsed: &'ast ParsedModuleRef) -> &'ast ast::Expr {
70        self._target(db).node(parsed)
71    }
72
73    /// Returns the scope where the unpack target expression belongs to.
74    pub fn target_scope(self, db: &'db dyn Db) -> ScopeId<'db> {
75        self.target_file_scope(db)
76            .to_scope_id(db, self.program_file(db))
77    }
78
79    pub fn program(self, db: &'db dyn Db) -> Program<'db> {
80        self.target_scope(db).program(db)
81    }
82
83    /// Returns the range of the unpack target expression.
84    pub fn range(self, db: &'db dyn Db, module: &ParsedModuleRef) -> TextRange {
85        self.target(db, module).range()
86    }
87}
88
89/// The expression that is being unpacked.
90#[derive(Clone, Copy, Debug, Hash, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
91pub struct UnpackValue<'db> {
92    /// The kind of unpack expression
93    kind: UnpackKind,
94    /// The expression we are unpacking
95    expression: Expression<'db>,
96}
97
98impl<'db> UnpackValue<'db> {
99    pub(crate) fn new(kind: UnpackKind, expression: Expression<'db>) -> Self {
100        Self { kind, expression }
101    }
102
103    /// Returns the underlying [`Expression`] that is being unpacked.
104    pub const fn expression(self) -> Expression<'db> {
105        self.expression
106    }
107
108    /// Returns the expression as an [`AnyNodeRef`].
109    pub fn as_any_node_ref<'ast>(
110        self,
111        db: &'db dyn Db,
112        module: &'ast ParsedModuleRef,
113    ) -> AnyNodeRef<'ast> {
114        self.expression().node_ref(db).node(module).into()
115    }
116
117    pub const fn kind(self) -> UnpackKind {
118        self.kind
119    }
120}
121
122#[derive(Clone, Copy, Debug, Hash, PartialEq, get_size2::GetSize)]
123pub enum UnpackKind {
124    /// An iterable expression like the one in a `for` loop or a comprehension.
125    Iterable { mode: EvaluationMode },
126    /// An context manager expression like the one in a `with` statement.
127    ContextManager { mode: EvaluationMode },
128    /// An expression that is being assigned to a target.
129    Assign,
130}
131
132/// The position of the target element in an unpacking.
133#[derive(Clone, Copy, Debug, Hash, PartialEq, get_size2::GetSize)]
134pub enum UnpackPosition {
135    /// The target element is in the first position of the unpacking.
136    First,
137    /// The target element is in the position other than the first position of the unpacking.
138    Other,
139}