Skip to content

Commit 178bb8e

Browse files
sharkdpGlyphack
authored andcommitted
[red-knot] Add custom __setattr__ support (astral-sh#16748)
## Summary Add support for classes with a custom `__setattr__` method. ## Test Plan New Markdown tests, ecosystem checks.
1 parent ba5ddaa commit 178bb8e

File tree

3 files changed

+119
-14
lines changed

3 files changed

+119
-14
lines changed

crates/red_knot_python_semantic/resources/mdtest/attributes.md

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1395,6 +1395,59 @@ def _(ns: argparse.Namespace):
13951395
reveal_type(ns.whatever) # revealed: Any
13961396
```
13971397

1398+
## Classes with custom `__setattr__` methods
1399+
1400+
### Basic
1401+
1402+
If a type provides a custom `__setattr__` method, we use the parameter type of that method as the
1403+
type to validate attribute assignments. Consider the following `CustomSetAttr` class:
1404+
1405+
```py
1406+
class CustomSetAttr:
1407+
def __setattr__(self, name: str, value: int) -> None:
1408+
pass
1409+
```
1410+
1411+
We can set arbitrary attributes on instances of this class:
1412+
1413+
```py
1414+
c = CustomSetAttr()
1415+
1416+
c.whatever = 42
1417+
```
1418+
1419+
### Type of the `name` parameter
1420+
1421+
If the `name` parameter of the `__setattr__` method is annotated with a (union of) literal type(s),
1422+
we only consider the attribute assignment to be valid if the assigned attribute is one of them:
1423+
1424+
```py
1425+
from typing import Literal
1426+
1427+
class Date:
1428+
def __setattr__(self, name: Literal["day", "month", "year"], value: int) -> None:
1429+
pass
1430+
1431+
date = Date()
1432+
date.day = 8
1433+
date.month = 4
1434+
date.year = 2025
1435+
1436+
# error: [unresolved-attribute] "Can not assign object of `Literal["UTC"]` to attribute `tz` on type `Date` with custom `__setattr__` method."
1437+
date.tz = "UTC"
1438+
```
1439+
1440+
### `argparse.Namespace`
1441+
1442+
A standard library example of a class with a custom `__setattr__` method is `argparse.Namespace`:
1443+
1444+
```py
1445+
import argparse
1446+
1447+
def _(ns: argparse.Namespace):
1448+
ns.whatever = 42
1449+
```
1450+
13981451
## Objects of all types have a `__class__` method
13991452

14001453
The type of `x.__class__` is the same as `x`'s meta-type. `x.__class__` is always the same value as

crates/red_knot_python_semantic/src/types.rs

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3422,13 +3422,31 @@ impl<'db> Type<'db> {
34223422
/// Returns an `Err` if the dunder method can't be called,
34233423
/// or the given arguments are not valid.
34243424
fn try_call_dunder(
3425+
self,
3426+
db: &'db dyn Db,
3427+
name: &str,
3428+
argument_types: CallArgumentTypes<'_, 'db>,
3429+
) -> Result<Bindings<'db>, CallDunderError<'db>> {
3430+
self.try_call_dunder_with_policy(db, name, argument_types, MemberLookupPolicy::empty())
3431+
}
3432+
3433+
/// Same as `try_call_dunder`, but allows specifying a policy for the member lookup. In
3434+
/// particular, this allows to specify `MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK` to avoid
3435+
/// looking up dunder methods on `object`, which is needed for functions like `__init__`,
3436+
/// `__new__`, or `__setattr__`.
3437+
fn try_call_dunder_with_policy(
34253438
self,
34263439
db: &'db dyn Db,
34273440
name: &str,
34283441
mut argument_types: CallArgumentTypes<'_, 'db>,
3442+
policy: MemberLookupPolicy,
34293443
) -> Result<Bindings<'db>, CallDunderError<'db>> {
34303444
match self
3431-
.member_lookup_with_policy(db, name.into(), MemberLookupPolicy::NO_INSTANCE_FALLBACK)
3445+
.member_lookup_with_policy(
3446+
db,
3447+
name.into(),
3448+
MemberLookupPolicy::NO_INSTANCE_FALLBACK | policy,
3449+
)
34323450
.symbol
34333451
{
34343452
Symbol::Type(dunder_callable, boundness) => {

crates/red_knot_python_semantic/src/types/infer.rs

Lines changed: 47 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,9 @@ use crate::types::{
8282
Truthiness, TupleType, Type, TypeAliasType, TypeAndQualifiers, TypeArrayDisplay,
8383
TypeQualifiers, TypeVarBoundOrConstraints, TypeVarInstance, UnionBuilder, UnionType,
8484
};
85-
use crate::types::{CallableType, FunctionDecorators, Signature};
85+
use crate::types::{
86+
CallableType, FunctionDecorators, MemberLookupPolicy, Signature, StringLiteralType,
87+
};
8688
use crate::unpack::{Unpack, UnpackPosition};
8789
use crate::util::subscript::{PyIndex, PySlice};
8890
use crate::Db;
@@ -2480,19 +2482,51 @@ impl<'db> TypeInferenceBuilder<'db> {
24802482

24812483
ensure_assignable_to(instance_attr_ty)
24822484
} else {
2483-
if emit_diagnostics {
2484-
self.context.report_lint(
2485-
&UNRESOLVED_ATTRIBUTE,
2486-
target,
2487-
format_args!(
2488-
"Unresolved attribute `{}` on type `{}`.",
2489-
attribute,
2490-
object_ty.display(db)
2491-
),
2492-
);
2493-
}
2485+
let result = object_ty.try_call_dunder_with_policy(
2486+
db,
2487+
"__setattr__",
2488+
CallArgumentTypes::positional([
2489+
Type::StringLiteral(StringLiteralType::new(
2490+
db,
2491+
Box::from(attribute),
2492+
)),
2493+
value_ty,
2494+
]),
2495+
MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK,
2496+
);
24942497

2495-
false
2498+
match result {
2499+
Ok(_) | Err(CallDunderError::PossiblyUnbound(_)) => true,
2500+
Err(CallDunderError::CallError(..)) => {
2501+
if emit_diagnostics {
2502+
self.context.report_lint(
2503+
&UNRESOLVED_ATTRIBUTE,
2504+
target,
2505+
format_args!(
2506+
"Can not assign object of `{}` to attribute `{attribute}` on type `{}` with custom `__setattr__` method.",
2507+
value_ty.display(db),
2508+
object_ty.display(db)
2509+
),
2510+
);
2511+
}
2512+
false
2513+
}
2514+
Err(CallDunderError::MethodNotAvailable) => {
2515+
if emit_diagnostics {
2516+
self.context.report_lint(
2517+
&UNRESOLVED_ATTRIBUTE,
2518+
target,
2519+
format_args!(
2520+
"Unresolved attribute `{}` on type `{}`.",
2521+
attribute,
2522+
object_ty.display(db)
2523+
),
2524+
);
2525+
}
2526+
2527+
false
2528+
}
2529+
}
24962530
}
24972531
}
24982532
}

0 commit comments

Comments
 (0)