Support dataclasses.InitVar

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/type_qualifiers/initvar.md

@@ -0,0 +1,124 @@
+# `dataclasses.InitVar`
+
+From the Python documentation on [`dataclasses.InitVar`]:
+
+If a field is an `InitVar`, it is considered a pseudo-field called an init-only field. As it is not
+a true field, it is not returned by the module-level `fields()` function. Init-only fields are added
+as parameters to the generated `__init__()` method, and are passed to the optional `__post_init__()`
+method. They are not otherwise used by dataclasses.
+
+## Basic
+
+Consider the following dataclass example where the `db` attribute is annotated with `InitVar`:
+
+```py
+from dataclasses import InitVar, dataclass
+
+class Database: ...
+
+@dataclass(order=True)
+class Person:
+    db: InitVar[Database]
+
+    name: str
+    age: int
+```
+
+We can see in the signature if `__init__`, that `db` is included as an argument:
+
+```py
+reveal_type(Person.__init__)  # revealed: (self: Person, db: Database, name: str, age: int) -> None
+```
+
+However, when we create an instance of this dataclass, the `db` attribute is not accessible:
+
+```py
+db = Database()
+alice = Person(db, "Alice", 30)
+
+alice.db  # error: [unresolved-attribute]
+```
+
+The `db` attribute is also not accessible on the class itself:
+
+```py
+Person.db  # error: [unresolved-attribute]
+```
+
+Other fields can still be accessed normally:
+
+```py
+reveal_type(alice.name)  # revealed: str
+reveal_type(alice.age)  # revealed: int
+```
+
+## `InitVar` wit default value
+
+An `InitVar` can also have a default value
+
+```py
+from dataclasses import InitVar, dataclass
+
+@dataclass
+class Person:
+    name: str
+    age: int
+
+    metadata: InitVar[str] = "default"
+
+reveal_type(Person.__init__)  # revealed: (self: Person, name: str, age: int, metadata: str = Literal["default"]) -> None
+
+alice = Person("Alice", 30)
+bob = Person("Bob", 25, "custom metadata")
+```
+
+## Error cases
+
+### Syntax
+
+`InitVar` can only be used with a single argument:
+
+```py
+from dataclasses import InitVar, dataclass
+
+@dataclass
+class Wrong:
+    x: InitVar[int, str]  # error: [invalid-type-form] "Type qualifier `InitVar` expected exactly 1 argument, got 2"
+```
+
+A bare `InitVar` is not allowed according to the [type annotation grammar]:
+
+```py
+@dataclass
+class AlsoWrong:
+    x: InitVar  # error: [invalid-type-form] "`InitVar` may not be used without a type argument"
+```
+
+### Outside of dataclasses
+
+`InitVar` annotations are not allowed outside of dataclass attribute annotations:
+
+```py
+from dataclasses import InitVar, dataclass
+
+# error: [invalid-type-form] "`InitVar` annotations are only allowed in class-body scopes"
+x: InitVar[int] = 1
+
+def f(x: InitVar[int]) -> None:  # error: [invalid-type-form] "`InitVar` is not allowed in function parameter annotations"
+    pass
+
+def g() -> InitVar[int]:  # error: [invalid-type-form] "`InitVar` is not allowed in function return type annotations"
+    return 1
+
+class C:
+    # TODO: this would ideally be an error
+    x: InitVar[int]
+
+@dataclass
+class D:
+    def __init__(self) -> None:
+        self.x: InitVar[int] = 1  # error: [invalid-type-form] "`InitVar` annotations are not allowed for non-name targets"
+```
+
+[type annotation grammar]: https://typing.python.org/en/latest/spec/annotations.html#type-and-annotation-expressions
+[`dataclasses.initvar`]: https://docs.python.org/3/library/dataclasses.html#dataclasses.InitVar

crates/ty_python_semantic/src/module_resolver/module.rs

@@ -344,10 +344,6 @@ impl KnownModule {
     pub const fn is_importlib(self) -> bool {
         matches!(self, Self::ImportLib)
     }
-
-    pub const fn is_dataclasses(self) -> bool {
-        matches!(self, Self::Dataclasses)
-    }
 }
 
 impl std::fmt::Display for KnownModule {

crates/ty_python_semantic/src/place.rs

@@ -524,6 +524,11 @@ impl<'db> PlaceAndQualifiers<'db> {
         self.qualifiers.contains(TypeQualifiers::CLASS_VAR)
     }
 
+    /// Returns `true` if the place has a `InitVar` type qualifier.
+    pub(crate) fn is_init_var(&self) -> bool {
+        self.qualifiers.contains(TypeQualifiers::INIT_VAR)
+    }
+
     /// Returns `Some(…)` if the place is qualified with `typing.Final` without a specified type.
     pub(crate) fn is_bare_final(&self) -> Option<TypeQualifiers> {
         match self {

crates/ty_python_semantic/src/types.rs

@@ -5214,7 +5214,7 @@ impl<'db> Type<'db> {
                     })
                 }
 
-                SpecialFormType::ClassVar | SpecialFormType::Final | SpecialFormType::InitVar => {
+                SpecialFormType::ClassVar | SpecialFormType::Final => {
                     Err(InvalidTypeExpressionError {
                         invalid_expressions: smallvec::smallvec_inline![
                             InvalidTypeExpression::TypeQualifier(*special_form)

crates/ty_python_semantic/src/types/class.rs

@@ -881,6 +881,13 @@ impl MethodDecorator {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct DataclassField<'db> {
+    pub(crate) field_ty: Type<'db>,
+    pub(crate) default_ty: Option<Type<'db>>,
+    pub(crate) init_only: bool,
+}
+
 /// Representation of a class definition statement in the AST: either a non-generic class, or a
 /// generic class that has not been specialized.
 ///
@@ -1473,6 +1480,10 @@ impl<'db> ClassLiteral<'db> {
             }
         }
 
+        if lookup_result.is_ok_and(|r| r.qualifiers.contains(TypeQualifiers::INIT_VAR)) {
+            return Place::Unbound.into();
+        }
+
         match (
             PlaceAndQualifiers::from(lookup_result),
             dynamic_type_to_intersect_with,
@@ -1580,10 +1591,16 @@ impl<'db> ClassLiteral<'db> {
 
         let signature_from_fields = |mut parameters: Vec<_>| {
             let mut kw_only_field_seen = false;
-            for (name, (mut attr_ty, mut default_ty)) in
-                self.fields(db, specialization, field_policy)
+            for (
+                field_name,
+                DataclassField {
+                    mut field_ty,
+                    mut default_ty,
+                    init_only: _,
+                },
+            ) in self.fields(db, specialization, field_policy)
             {
-                if attr_ty
+                if field_ty
                     .into_nominal_instance()
                     .is_some_and(|instance| instance.class.is_known(db, KnownClass::KwOnly))
                 {
@@ -1594,7 +1611,7 @@ impl<'db> ClassLiteral<'db> {
                     continue;
                 }
 
-                let dunder_set = attr_ty.class_member(db, "__set__".into());
+                let dunder_set = field_ty.class_member(db, "__set__".into());
                 if let Place::Type(dunder_set, Boundness::Bound) = dunder_set.place {
                     // The descriptor handling below is guarded by this not-dynamic check, because
                     // dynamic types like `Any` are valid (data) descriptors: since they have all
@@ -1623,7 +1640,7 @@ impl<'db> ClassLiteral<'db> {
                                 }
                             }
                         }
-                        attr_ty = value_types.build();
+                        field_ty = value_types.build();
 
                         // The default value of the attribute is *not* determined by the right hand side
                         // of the class-body assignment. Instead, the runtime invokes `__get__` on the
@@ -1640,11 +1657,11 @@ impl<'db> ClassLiteral<'db> {
                 }
 
                 let mut parameter = if kw_only_field_seen {
-                    Parameter::keyword_only(name)
+                    Parameter::keyword_only(field_name)
                 } else {
-                    Parameter::positional_or_keyword(name)
+                    Parameter::positional_or_keyword(field_name)
                 }
-                .with_annotated_type(attr_ty);
+                .with_annotated_type(field_ty);
 
                 if let Some(default_ty) = default_ty {
                     parameter = parameter.with_default_type(default_ty);
@@ -1746,7 +1763,7 @@ impl<'db> ClassLiteral<'db> {
         db: &'db dyn Db,
         specialization: Option<Specialization<'db>>,
         field_policy: CodeGeneratorKind,
-    ) -> FxOrderMap<Name, (Type<'db>, Option<Type<'db>>)> {
+    ) -> FxOrderMap<Name, DataclassField<'db>> {
         if field_policy == CodeGeneratorKind::NamedTuple {
             // NamedTuples do not allow multiple inheritance, so it is sufficient to enumerate the
             // fields of this class only.
@@ -1793,7 +1810,7 @@ impl<'db> ClassLiteral<'db> {
         self,
         db: &'db dyn Db,
         specialization: Option<Specialization<'db>>,
-    ) -> FxOrderMap<Name, (Type<'db>, Option<Type<'db>>)> {
+    ) -> FxOrderMap<Name, DataclassField<'db>> {
         let mut attributes = FxOrderMap::default();
 
         let class_body_scope = self.body_scope(db);
@@ -1835,11 +1852,12 @@ impl<'db> ClassLiteral<'db> {
 
                     attributes.insert(
                         place_expr.expect_name().clone(),
-                        (
-                            attr_ty.apply_optional_specialization(db, specialization),
-                            default_ty
+                        DataclassField {
+                            field_ty: attr_ty.apply_optional_specialization(db, specialization),
+                            default_ty: default_ty
                                 .map(|ty| ty.apply_optional_specialization(db, specialization)),
-                        ),
+                            init_only: attr.is_init_var(),
+                        },
                     );
                 }
             }
@@ -1876,6 +1894,10 @@ impl<'db> ClassLiteral<'db> {
                         qualifiers,
                     } = class.own_instance_member(db, name)
                     {
+                        if qualifiers.contains(TypeQualifiers::INIT_VAR) {
+                            continue;
+                        }
+
                         // TODO: We could raise a diagnostic here if there are conflicting type qualifiers
                         union_qualifiers |= qualifiers;
 
@@ -2592,6 +2614,7 @@ pub enum KnownClass {
     // dataclasses
     Field,
     KwOnly,
+    InitVar,
     // _typeshed._type_checker_internals
     NamedTupleFallback,
 }
@@ -2686,6 +2709,7 @@ impl KnownClass {
             | Self::Deprecated
             | Self::Field
             | Self::KwOnly
+            | Self::InitVar
             | Self::NamedTupleFallback => Truthiness::Ambiguous,
         }
     }
@@ -2744,6 +2768,7 @@ impl KnownClass {
             | Self::EllipsisType
             | Self::NotImplementedType
             | Self::KwOnly
+            | Self::InitVar
             | Self::VersionInfo
             | Self::Bool
             | Self::NoneType => false,
@@ -2843,6 +2868,7 @@ impl KnownClass {
             | KnownClass::NotImplementedType
             | KnownClass::Field
             | KnownClass::KwOnly
+            | KnownClass::InitVar
             | KnownClass::NamedTupleFallback => false,
         }
     }
@@ -2925,6 +2951,7 @@ impl KnownClass {
             | Self::UnionType
             | Self::Field
             | Self::KwOnly
+            | Self::InitVar
             | Self::NamedTupleFallback => false,
         }
     }
@@ -3016,6 +3043,7 @@ impl KnownClass {
             Self::NotImplementedType => "_NotImplementedType",
             Self::Field => "Field",
             Self::KwOnly => "KW_ONLY",
+            Self::InitVar => "InitVar",
             Self::NamedTupleFallback => "NamedTupleFallback",
         }
     }
@@ -3269,8 +3297,7 @@ impl KnownClass {
             | Self::DefaultDict
             | Self::Deque
             | Self::OrderedDict => KnownModule::Collections,
-            Self::Field => KnownModule::Dataclasses,
-            Self::KwOnly => KnownModule::Dataclasses,
+            Self::Field | Self::KwOnly | Self::InitVar => KnownModule::Dataclasses,
             Self::NamedTupleFallback => KnownModule::TypeCheckerInternals,
         }
     }
@@ -3342,6 +3369,7 @@ impl KnownClass {
             | Self::NewType
             | Self::Field
             | Self::KwOnly
+            | Self::InitVar
             | Self::Iterable
             | Self::Iterator
             | Self::NamedTupleFallback => false,
@@ -3417,6 +3445,7 @@ impl KnownClass {
             | Self::NewType
             | Self::Field
             | Self::KwOnly
+            | Self::InitVar
             | Self::Iterable
             | Self::Iterator
             | Self::NamedTupleFallback => false,
@@ -3504,6 +3533,7 @@ impl KnownClass {
             "_NotImplementedType" => Self::NotImplementedType,
             "Field" => Self::Field,
             "KW_ONLY" => Self::KwOnly,
+            "InitVar" => Self::InitVar,
             "NamedTupleFallback" => Self::NamedTupleFallback,
             _ => return None,
         };
@@ -3566,6 +3596,7 @@ impl KnownClass {
             | Self::WrapperDescriptorType
             | Self::Field
             | Self::KwOnly
+            | Self::InitVar
             | Self::NamedTupleFallback => module == self.canonical_module(db),
             Self::NoneType => matches!(module, KnownModule::Typeshed | KnownModule::Types),
             Self::SpecialForm

crates/ty_python_semantic/src/types/class_base.rs

@@ -192,8 +192,7 @@ impl<'db> ClassBase<'db> {
                 | SpecialFormType::TypeOf
                 | SpecialFormType::CallableTypeOf
                 | SpecialFormType::AlwaysTruthy
-                | SpecialFormType::AlwaysFalsy
-                | SpecialFormType::InitVar => None,
+                | SpecialFormType::AlwaysFalsy => None,
 
                 SpecialFormType::Unknown => Some(Self::unknown()),