diff --git a/crates/pyrefly_types/src/function.rs b/crates/pyrefly_types/src/function.rs index b7e88ebc5a..cb7d22e338 100644 --- a/crates/pyrefly_types/src/function.rs +++ b/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, } } diff --git a/pyrefly/lib/alt/call.rs b/pyrefly/lib/alt/call.rs index d9624854f4..dfafec80cf 100644 --- a/pyrefly/lib/alt/call.rs +++ b/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 diff --git a/pyrefly/lib/alt/class/class_field.rs b/pyrefly/lib/alt/class/class_field.rs index e89f5105d0..6951861ea8 100644 --- a/pyrefly/lib/alt/class/class_field.rs +++ b/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, diff --git a/pyrefly/lib/alt/special_calls.rs b/pyrefly/lib/alt/special_calls.rs index 054498dc51..878962c99d 100644 --- a/pyrefly/lib/alt/special_calls.rs +++ b/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, + 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 { + // 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 { + match expr { + Expr::BooleanLiteral(value) => Some(value.value), + _ => None, + } + } + + fn python_type_from_type_engine_expr(&self, expr: &Expr) -> Option { + 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 { + 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 { + 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, diff --git a/pyrefly/lib/export/special.rs b/pyrefly/lib/export/special.rs index 89bc8d331d..fbe84313cd 100644 --- a/pyrefly/lib/export/special.rs +++ b/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"), diff --git a/pyrefly/lib/test/mod.rs b/pyrefly/lib/test/mod.rs index f25778ea0a..b5a7aa0ca8 100644 --- a/pyrefly/lib/test/mod.rs +++ b/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; diff --git a/pyrefly/lib/test/sqlalchemy.rs b/pyrefly/lib/test/sqlalchemy.rs new file mode 100644 index 0000000000..9fade5bf31 --- /dev/null +++ b/pyrefly/lib/test/sqlalchemy.rs @@ -0,0 +1,201 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +use crate::test::util::TestEnv; +use crate::testcase; + +fn sqlalchemy_env() -> TestEnv { + let mut env = TestEnv::new(); + env.add_with_path( + "sqlalchemy", + "sqlalchemy/__init__.py", + "from . import orm\n", + ); + env.add_with_path( + "sqlalchemy.sql", + "sqlalchemy/sql/__init__.py", + "from . import sqltypes\nfrom . import type_api\n", + ); + env.add_with_path( + "sqlalchemy.sql.type_api", + "sqlalchemy/sql/type_api.py", + r#" +from typing import Generic, TypeVar + +_T = TypeVar("_T") + +class TypeEngine(Generic[_T]): + ... + +class TypeDecorator(TypeEngine[_T], Generic[_T]): + ... +"#, + ); + env.add_with_path( + "sqlalchemy.sql.sqltypes", + "sqlalchemy/sql/sqltypes.py", + r#" +from enum import Enum as PythonEnum +from typing import Generic, TypeVar +from .type_api import TypeEngine + +_T = TypeVar("_T") + +class String(TypeEngine[str]): + ... + +class Integer(TypeEngine[int]): + ... + +class Float(TypeEngine[float]): + ... + +class GenericType(TypeEngine[_T], Generic[_T]): + ... + +class Enum(TypeEngine[str | PythonEnum]): + def __init__(self, enum_class: type[PythonEnum]) -> None: + ... +"#, + ); + env.add_with_path( + "sqlalchemy.orm.base", + "sqlalchemy/orm/base.py", + r#" +from typing import Generic, TypeVar + +_T = TypeVar("_T") + +class Mapped(Generic[_T]): + ... + +class _MappedAnnotationBase(Mapped[_T]): + ... + +class _DeclarativeMapped(_MappedAnnotationBase[_T]): + ... +"#, + ); + env.add_with_path( + "sqlalchemy.orm.properties", + "sqlalchemy/orm/properties.py", + r#" +from typing import Generic, TypeVar +from .base import _DeclarativeMapped + +_T = TypeVar("_T") + +class MappedColumn(_DeclarativeMapped[_T]): + column: object + + def __init__(self) -> None: + ... +"#, + ); + env.add_with_path( + "sqlalchemy.orm", + "sqlalchemy/orm/__init__.py", + r#" +from typing import Any +from sqlalchemy.sql.type_api import TypeEngine +from .base import Mapped +from .properties import MappedColumn + +__all__ = ["Mapped", "MappedColumn", "mapped_column"] + +def mapped_column( + type_: TypeEngine[Any] | type[TypeEngine[Any]] | None = None, + *args: Any, + **kw: Any, +) -> MappedColumn[Any]: + return MappedColumn() +"#, + ); + env +} + +testcase!( + test_sqlalchemy_mapped_column_infers_type, + sqlalchemy_env(), + r#" +from typing import assert_type +from sqlalchemy.orm import MappedColumn, mapped_column +from sqlalchemy.sql.sqltypes import Integer, String + +class Model: + name = mapped_column(String()) + quantity = mapped_column(Integer) + sku = mapped_column(String(), nullable=False) + pk = mapped_column(Integer, primary_key=True) + nullable_pk = mapped_column(Integer, primary_key=True, nullable=True) + +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]) +assert_type(Model.nullable_pk, MappedColumn[int | None]) + +def make_column() -> MappedColumn[str | None]: + result = mapped_column(String()) + result.column + return result +"#, +); + +testcase!( + test_sqlalchemy_mapped_column_substitutes_generic_types, + sqlalchemy_env(), + r#" +from typing import assert_type, reveal_type +from sqlalchemy.orm import MappedColumn, mapped_column +from sqlalchemy.sql.sqltypes import GenericType +from sqlalchemy.sql.type_api import TypeDecorator + +assert_type(mapped_column(GenericType[float]()), MappedColumn[float | None]) +assert_type(mapped_column(TypeDecorator[str]()), MappedColumn[str | None]) +reveal_type(mapped_column(GenericType)) # E: revealed type: MappedColumn[Unknown | None] +"#, +); + +testcase!( + test_sqlalchemy_mapped_column_infers_enum_class, + sqlalchemy_env(), + r#" +from enum import Enum as PythonEnum +from typing import assert_type +from sqlalchemy.orm import MappedColumn, mapped_column +from sqlalchemy.sql.sqltypes import Enum + +class Color(PythonEnum): + RED = "red" + +assert_type(mapped_column(Enum(Color)), MappedColumn[Color | None]) +"#, +); + +testcase!( + test_sqlalchemy_mapped_column_respects_explicit_annotation, + sqlalchemy_env(), + r#" +from typing import Literal, assert_type +from sqlalchemy.orm import Mapped, mapped_column +from sqlalchemy.sql.sqltypes import Float, Integer, String + +Key = Literal["api_key", "api_secret"] + +class Model: + key: Mapped[Key] = mapped_column(String(), nullable=False) + quantity: Mapped[int] = mapped_column(Integer, primary_key=True) + ratio: Mapped[float] = mapped_column(Float(), nullable=False) + maybe_ratio: Mapped[float | None] = mapped_column(Float()) + +assert_type(Model.key, Mapped[Key]) +assert_type(Model.quantity, Mapped[int]) +assert_type(Model.ratio, Mapped[float]) +assert_type(Model.maybe_ratio, Mapped[float | None]) +"#, +);