fix Infer SQLAlchemy mapped_column type #1600

crates/pyrefly_types/src/function.rs

@@ -473,6 +473,8 @@ pub enum FunctionKind {
     UsesShapeDsl,
     /// The `shape_extensions.defines_assert_shape` decorator function itself.
     DefinesAssertShape,
+    /// `sqlalchemy.orm.mapped_column()`.
+    SqlAlchemyMappedColumn,
 }
 
 impl FunctionKind {
@@ -536,6 +538,9 @@ impl FunctionKind {
             ("numba.core.decorators", "njit") => Self::NumbaNjit,
             ("shape_extensions", "uses_shape_dsl") => Self::UsesShapeDsl,
             ("shape_extensions", "defines_assert_shape") => Self::DefinesAssertShape,
+            ("sqlalchemy.orm" | "sqlalchemy.orm._orm_constructors", "mapped_column") => {
+                Self::SqlAlchemyMappedColumn
+            }
             _ => Self::Def(id),
         }
     }
@@ -578,6 +583,7 @@ impl FunctionKind {
             Self::ShapeDsl(id, _, _) | Self::TypeShapeDsl(id, _) => id.qname.module_name(),
             Self::UsesShapeDsl => ModuleName::from_str("shape_extensions"),
             Self::DefinesAssertShape => ModuleName::from_str("shape_extensions"),
+            Self::SqlAlchemyMappedColumn => ModuleName::from_str("sqlalchemy.orm"),
         }
     }
 
@@ -619,6 +625,7 @@ impl FunctionKind {
             Self::ShapeDsl(id, _, _) | Self::TypeShapeDsl(id, _) => Cow::Borrowed(id.qname.id()),
             Self::UsesShapeDsl => Cow::Owned(Name::new_static("uses_shape_dsl")),
             Self::DefinesAssertShape => Cow::Owned(Name::new_static("defines_assert_shape")),
+            Self::SqlAlchemyMappedColumn => Cow::Owned(Name::new_static("mapped_column")),
         }
     }
 
@@ -660,6 +667,7 @@ impl FunctionKind {
             Self::ShapeDsl(id, _, _) | Self::TypeShapeDsl(id, _) => id.cls.clone(),
             Self::UsesShapeDsl => None,
             Self::DefinesAssertShape => None,
+            Self::SqlAlchemyMappedColumn => None,
         }
     }
 

pyrefly/lib/alt/call.rs

@@ -3147,6 +3147,17 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                         errors,
                     )
                 }
+                Some(CalleeKind::Function(FunctionKind::SqlAlchemyMappedColumn)) => self
+                    .call_sqlalchemy_mapped_column(
+                        ty.clone(),
+                        &args,
+                        &kws,
+                        x.func.range(),
+                        x.arguments.range,
+                        x,
+                        hint,
+                        errors,
+                    ),
                 // Treat assert_type and reveal_type like pseudo-builtins for convenience. Note that we still
                 // log a name-not-found error, but we also assert/reveal the type as requested.
                 None if ty.is_error() && is_special_name(&x.func, "assert_type") => self

pyrefly/lib/alt/class/class_field.rs

@@ -1584,10 +1584,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         // container defaults.
         if !has_keyword_conflict
             && let Some(default) = mutable_default
-            && function_id.is_some_and(|id| {
-                id.has_toplevel_qname("sqlalchemy.orm", "mapped_column")
-                    || id.has_toplevel_qname("sqlalchemy.orm._orm_constructors", "mapped_column")
-            })
+            && matches!(
+                callee_kind,
+                Some(CalleeKind::Function(FunctionKind::SqlAlchemyMappedColumn))
+            )
         {
             self.error(
                 errors,

pyrefly/lib/alt/special_calls.rs

@@ -15,9 +15,11 @@ use pyrefly_python::dunder;
 use pyrefly_types::shaped_array::IntTuple;
 use pyrefly_types::shaped_array::IntTupleView;
 use pyrefly_types::shaped_array::is_tuple_carrier_shape_middle;
+use pyrefly_types::type_alias::TypeAliasData;
 use pyrefly_util::visit::Visit;
 use pyrefly_util::visit::VisitMut;
 use ruff_python_ast::Expr;
+use ruff_python_ast::ExprCall;
 use ruff_python_ast::Keyword;
 use ruff_python_ast::name::Name;
 use ruff_text_size::Ranged;
@@ -40,6 +42,7 @@ use crate::types::callable::Param;
 use crate::types::callable::Params;
 use crate::types::callable::unexpected_keyword;
 use crate::types::class::Class;
+use crate::types::class::ClassType;
 use crate::types::function::FunctionKind;
 use crate::types::tuple::Tuple;
 use crate::types::types::Forallable;
@@ -525,6 +528,186 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         ret
     }
 
+    pub fn call_sqlalchemy_mapped_column(
+        &self,
+        callee_ty: Type,
+        args: &[CallArg],
+        keywords: &[CallKeyword],
+        callee_range: TextRange,
+        arg_range: TextRange,
+        call: &ExprCall,
+        hint: Option<HintRef>,
+        errors: &ErrorCollector,
+    ) -> Type {
+        let ret = self
+            .freeform_call_infer(
+                callee_ty,
+                args,
+                keywords,
+                callee_range,
+                arg_range,
+                hint,
+                errors,
+            )
+            .ty;
+        if hint.is_some_and(|hint| {
+            hint.types()
+                .iter()
+                .any(|ty| self.is_sqlalchemy_mapped_hint(ty))
+        }) {
+            return ret;
+        }
+        let Some(mut python_type) = self.sqlalchemy_mapped_column_python_type(call) else {
+            return ret;
+        };
+        if self.sqlalchemy_mapped_column_is_nullable(call) {
+            python_type = Type::optional(python_type);
+        }
+        self.apply_sqlalchemy_mapped_python_type(ret, python_type)
+    }
+
+    fn sqlalchemy_mapped_column_python_type(&self, call: &ExprCall) -> Option<Type> {
+        // Check up to the first two positional arguments for a SQLAlchemy TypeEngine-derived column type.
+        for expr in call.arguments.args.iter().take(2) {
+            if let Some(ty) = self.python_type_from_type_engine_expr(expr) {
+                return Some(ty);
+            }
+        }
+        for keyword in &call.arguments.keywords {
+            if let Some(arg) = &keyword.arg
+                && matches!(arg.as_str(), "type_")
+                && let Some(ty) = self.python_type_from_type_engine_expr(&keyword.value)
+            {
+                return Some(ty);
+            }
+        }
+        None
+    }
+
+    fn sqlalchemy_mapped_column_is_nullable(&self, call: &ExprCall) -> bool {
+        let mut nullable = None;
+        let mut primary_key = false;
+        for keyword in &call.arguments.keywords {
+            let Some(arg) = &keyword.arg else {
+                continue;
+            };
+            match arg.as_str() {
+                "nullable" => {
+                    if let Some(value) = Self::expr_bool_literal(&keyword.value) {
+                        nullable = Some(value);
+                    }
+                }
+                "primary_key" => {
+                    if let Some(value) = Self::expr_bool_literal(&keyword.value) {
+                        primary_key = value;
+                    }
+                }
+                _ => {}
+            }
+        }
+        nullable.unwrap_or(!primary_key)
+    }
+
+    fn expr_bool_literal(expr: &Expr) -> Option<bool> {
+        match expr {
+            Expr::BooleanLiteral(value) => Some(value.value),
+            _ => None,
+        }
+    }
+
+    fn python_type_from_type_engine_expr(&self, expr: &Expr) -> Option<Type> {
+        if let Expr::Call(call) = expr
+            && matches!(
+                self.expr_infer(&call.func, &self.error_swallower()),
+                Type::ClassDef(cls)
+                    if cls.has_toplevel_qname("sqlalchemy.sql.sqltypes", "Enum")
+            )
+            && let Some(enum_class) = call.arguments.args.first()
+        {
+            let enum_class = self.expr_infer(enum_class, &self.error_swallower());
+            return match enum_class {
+                Type::ClassDef(cls) => Some(self.promote_silently(&cls)),
+                Type::Type(instance) => Some(*instance),
+                _ => None,
+            };
+        }
+        let ty = self.expr_infer(expr, &self.error_swallower());
+        self.python_type_from_type_engine_type(&ty)
+    }
+
+    fn python_type_from_type_engine_type(&self, ty: &Type) -> Option<Type> {
+        match ty {
+            Type::ClassType(cls) => self.python_type_from_type_engine_class(cls),
+            Type::ClassDef(cls) => {
+                let inst = self.promote_silently(cls);
+                self.python_type_from_type_engine_type(&inst)
+            }
+            Type::Type(inner) => self.python_type_from_type_engine_type(inner),
+            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 = Vec::new();
+                for member in &u.members {
+                    if let Some(member_ty) = self.python_type_from_type_engine_type(member) {
+                        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))
+                    }
+                }
+            }
+            _ => None,
+        }
+    }
+
+    fn python_type_from_type_engine_class(&self, cls: &ClassType) -> Option<Type> {
+        let type_engine = self
+            .get_mro_for_class(cls.class_object())
+            .ancestors(self.stdlib)
+            .find(|ancestor| {
+                ancestor
+                    .class_object()
+                    .has_toplevel_qname("sqlalchemy.sql.type_api", "TypeEngine")
+            })?;
+        self.as_superclass(cls, type_engine.class_object())?
+            .targs()
+            .as_slice()
+            .first()
+            .cloned()
+    }
+
+    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
+                && class_type
+                    .class_object()
+                    .has_toplevel_qname("sqlalchemy.orm.properties", "MappedColumn")
+                && let Some(slot) = class_type.targs_mut().as_mut().get_mut(0)
+            {
+                *slot = python_type.clone();
+            }
+        });
+        ty
+    }
+
+    fn is_sqlalchemy_mapped_hint(&self, ty: &Type) -> bool {
+        matches!(ty, Type::ClassType(cls) if cls.has_qname("sqlalchemy.orm.base", "Mapped"))
+    }
+
     pub fn call_isinstance(
         &self,
         obj: &Expr,

pyrefly/lib/export/special.rs

@@ -76,6 +76,7 @@ pub enum SpecialExport {
     BuiltinsFrozenset,
     BuiltinsFloat,
     Deprecated,
+    SqlAlchemyMappedColumn,
     Final,
     TypingMapping,
     TypeForm,
@@ -163,6 +164,7 @@ impl SpecialExport {
             "frozenset" => Some(Self::BuiltinsFrozenset),
             "float" => Some(Self::BuiltinsFloat),
             "deprecated" => Some(Self::Deprecated),
+            "mapped_column" => Some(Self::SqlAlchemyMappedColumn),
             "Final" => Some(Self::Final),
             "Mapping" => Some(Self::TypingMapping),
             "TypeForm" => Some(Self::TypeForm),
@@ -268,6 +270,12 @@ impl SpecialExport {
             Self::ShapeVars => matches!(m.as_str(), "shape_extensions"),
             Self::ProxyMethod => matches!(m.as_str(), "shape_extensions"),
             Self::Sentinel => matches!(m.as_str(), "typing_extensions"),
+            Self::SqlAlchemyMappedColumn => {
+                matches!(
+                    m.as_str(),
+                    "sqlalchemy.orm" | "sqlalchemy.orm._orm_constructors"
+                )
+            }
             // `builtins.sentinel` (3.15+) and its `typing_extensions.sentinel`
             // backport are the same lowercase PEP 661 constructor.
             Self::BuiltinsSentinel => matches!(m.as_str(), "builtins" | "typing_extensions"),

pyrefly/lib/test/mod.rs

@@ -89,6 +89,7 @@ mod shaped_comments;
 mod simple;
 mod sklearn;
 mod slots;
+mod sqlalchemy;
 mod state;
 mod subscript_narrow;
 mod suppression;