.

.

.

Author: asukaminato0721

crates/pyrefly_types/src/class.rs

@@ -249,9 +249,7 @@ impl ClassKind {
             ("enum", "property") => Self::Property(name.clone()),
             ("enum", "member") => Self::EnumMember,
             ("enum", "nonmember") => Self::EnumNonmember,
-            ("dataclasses", "Field")
-            | ("sqlalchemy.orm", "MappedColumn")
-            | ("sqlalchemy.orm.properties", "MappedColumn") => Self::DataclassField,
+            ("dataclasses", "Field") => Self::DataclassField,
             _ => Self::Class,
         }
     }

pyrefly/lib/alt/special_calls.rs

@@ -12,6 +12,7 @@
  */
 
 use pyrefly_types::callable::FuncMetadata;
+use pyrefly_types::type_alias::TypeAliasData;
 use pyrefly_types::types::Union;
 use pyrefly_util::visit::Visit;
 use pyrefly_util::visit::VisitMut;
@@ -39,8 +40,6 @@ use crate::types::callable::FunctionKind;
 use crate::types::callable::unexpected_keyword;
 use crate::types::class::Class;
 use crate::types::class::ClassType;
-use crate::types::special_form::SpecialForm;
-use crate::types::special_form::SpecialForm;
 use crate::types::tuple::Tuple;
 use crate::types::types::Type;
 
@@ -338,7 +337,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         }
         for keyword in &call.arguments.keywords {
             if let Some(arg) = &keyword.arg
-                && matches!(arg.as_str(), "__type_pos" | "type_")
+                && matches!(arg.as_str(), "type_")
                 && let Some(ty) = self.python_type_from_type_engine_expr(&keyword.value)
             {
                 return Some(ty);
@@ -368,7 +367,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 _ => {}
             }
         }
-        if primary_key { false } else { nullable }
+        !primary_key && nullable
     }
 
     fn expr_bool_literal(expr: &Expr) -> Option<bool> {
@@ -391,61 +390,70 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 self.python_type_from_type_engine_type(&inst)
             }
             Type::Type(inner) => self.python_type_from_type_engine_type(inner),
-            Type::TypeAlias(alias) => self.python_type_from_type_engine_type(&alias.as_type()),
+            Type::TypeAlias(alias) => match alias.as_ref() {
+                TypeAliasData::Value(alias) => {
+                    self.python_type_from_type_engine_type(&alias.as_type())
+                }
+                TypeAliasData::Ref(_) => {
+                    self.python_type_from_type_engine_type(&self.untype_alias(alias))
+                }
+            },
             Type::Union(u) => {
-                let mut inferred = None;
+                let mut inferred = Vec::new();
                 for member in &u.members {
                     if let Some(member_ty) = self.python_type_from_type_engine_type(member) {
-                        match &inferred {
-                            Some(existing) if existing != &member_ty => return None,
-                            None => inferred = Some(member_ty),
-                            _ => {}
-                        }
+                        inferred.push(member_ty);
+                    }
+                }
+                if inferred.is_empty() {
+                    None
+                } else {
+                    inferred.sort();
+                    inferred.dedup();
+                    if inferred.len() == 1 {
+                        inferred.pop()
+                    } else {
+                        Some(Type::union(inferred))
                     }
                 }
-                inferred
             }
             _ => None,
         }
     }
 
     fn python_type_from_type_engine_class(&self, cls: &ClassType) -> Option<Type> {
-        if Self::is_sqlalchemy_type_engine_class(cls.class_object()) {
+        if cls
+            .class_object()
+            .has_toplevel_qname("sqlalchemy.sql.type_api", "TypeEngine")
+        {
             return cls.targs().as_slice().first().cloned();
         }
-        let bases = self.get_base_types_for_class(cls.class_object());
-        for base in bases.iter() {
-            if let Some(ty) = self.python_type_from_type_engine_class(base) {
-                return Some(ty);
+        let mro = self.get_mro_for_class(cls.class_object());
+        for ancestor in mro.ancestors_no_object() {
+            if ancestor
+                .class_object()
+                .has_toplevel_qname("sqlalchemy.sql.type_api", "TypeEngine")
+            {
+                return ancestor.targs().as_slice().first().cloned();
             }
         }
         None
     }
 
-    fn is_sqlalchemy_type_engine_class(class: &Class) -> bool {
-        class.has_toplevel_qname("sqlalchemy.sql.type_api", "TypeEngine")
-    }
-
     fn apply_sqlalchemy_mapped_python_type(&self, mut ty: Type, python_type: Type) -> Type {
         ty.visit_mut(&mut |inner| {
             if let Type::ClassType(class_type) = inner
-                && Self::is_sqlalchemy_mapped_class(class_type.class_object())
+                && class_type
+                    .class_object()
+                    .has_toplevel_qname("sqlalchemy.orm.properties", "MappedColumn")
+                && let Some(slot) = class_type.targs_mut().as_mut().get_mut(0)
             {
-                if let Some(slot) = class_type.targs_mut().as_mut().get_mut(0) {
-                    *slot = python_type.clone();
-                }
+                *slot = python_type.clone();
             }
         });
         ty
     }
 
-    fn is_sqlalchemy_mapped_class(class: &Class) -> bool {
-        class.has_toplevel_qname("sqlalchemy.orm.base", "Mapped")
-            || class.has_toplevel_qname("sqlalchemy.orm.base", "_MappedAnnotationBase")
-            || class.has_toplevel_qname("sqlalchemy.orm.base", "_DeclarativeMapped")
-            || class.has_toplevel_qname("sqlalchemy.orm.properties", "MappedColumn")
-    }
-
     pub fn call_isinstance(
         &self,
         obj: &Expr,

pyrefly/lib/export/special.rs

@@ -205,13 +205,8 @@ impl SpecialExport {
                 m.as_str(),
                 "typing" | "typing_extensions" | "collections.abc"
             ),
+            Self::SqlAlchemyMappedColumn => matches!(m.as_str(), "sqlalchemy.orm"),
             Self::Deprecated => matches!(m.as_str(), "warnings" | "typing_extensions"),
-            Self::SqlAlchemyMappedColumn => {
-                matches!(
-                    m.as_str(),
-                    "sqlalchemy.orm" | "sqlalchemy.orm._orm_constructors"
-                )
-            }
         }
     }
 

pyrefly/lib/test/mod.rs

@@ -65,8 +65,8 @@ mod returns;
 mod scope;
 mod semantic_syntax_errors;
 mod simple;
-mod sqlalchemy;
 mod slots;
+mod sqlalchemy;
 mod state;
 mod subscript_narrow;
 mod suppression;

pyrefly/lib/test/sqlalchemy.rs

@@ -102,8 +102,8 @@ testcase!(
     test_sqlalchemy_mapped_column_infers_type,
     sqlalchemy_env(),
     r#"
-from typing import reveal_type
-from sqlalchemy.orm import mapped_column
+from typing import assert_type
+from sqlalchemy.orm import MappedColumn, mapped_column
 from sqlalchemy.sql.sqltypes import Integer, String
 
 class Model:
@@ -112,9 +112,9 @@ class Model:
     sku = mapped_column(String(), nullable=False)
     pk = mapped_column(Integer, primary_key=True)
 
-reveal_type(Model.name)  # E: revealed type: MappedColumn[str | None]
-reveal_type(Model.quantity)  # E: revealed type: MappedColumn[int | None]
-reveal_type(Model.sku)  # E: revealed type: MappedColumn[str]
-reveal_type(Model.pk)  # E: revealed type: MappedColumn[int]
+assert_type(Model.name, MappedColumn[str | None])
+assert_type(Model.quantity, MappedColumn[int | None])
+assert_type(Model.sku, MappedColumn[str])
+assert_type(Model.pk, MappedColumn[int])
 "#,
 );