feat: enhance deferred_column_property with additional parameters and typing tests for compatibility

Author: 50Bytes-dev

sqlmodel/deferred_column.py

@@ -5,16 +5,34 @@
 with safe deferred loading - returning fallback values instead of raising DetachedInstanceError.
 """
 
-from typing import Any
+from typing import TYPE_CHECKING, Any, Optional, Type, TypeVar
 
 from sqlalchemy import event
-from sqlalchemy.orm import ColumnProperty, column_property
+from sqlalchemy.orm import column_property
 from sqlalchemy.orm.attributes import set_committed_value
 
+if TYPE_CHECKING:
+    from sqlalchemy.orm._typing import _ORMColumnExprArgument
+    from sqlalchemy.orm.interfaces import PropComparator
+    from sqlalchemy.orm.properties import MappedSQLExpression
+    from sqlalchemy.sql._typing import _InfoType
+
+_T = TypeVar("_T")
+
 
 def deferred_column_property(
-    expression: Any, *, fallback_value: Any, deferred: bool = True, **kwargs: Any
-) -> ColumnProperty:
+    expression: "_ORMColumnExprArgument[_T]",
+    *additional_expressions: "_ORMColumnExprArgument[Any]",
+    fallback_value: Any,
+    group: Optional[str] = None,
+    deferred: bool = True,
+    raiseload: bool = False,
+    comparator_factory: Optional[Type["PropComparator[_T]"]] = None,
+    active_history: bool = False,
+    expire_on_flush: bool = True,
+    info: Optional["_InfoType"] = None,
+    doc: Optional[str] = None,
+) -> "MappedSQLExpression[_T]":
     """
     Create a deferred column property that returns a fallback value instead of raising
     DetachedInstanceError when accessed on a detached instance.
@@ -24,12 +42,19 @@ def deferred_column_property(
 
     Args:
         expression: The SQL expression for the column property
+        *additional_expressions: Additional SQL expressions for the column property
         fallback_value: Value to return when property cannot be loaded (required)
+        group: A group name for this property when marked as deferred
         deferred: Whether to defer loading the property (defaults to True)
-        **kwargs: Additional arguments passed to column_property
+        raiseload: When True, loading this property will raise an error
+        comparator_factory: A class which extends PropComparator for custom SQL clause generation
+        active_history: When True, indicates that the "previous" value should be loaded when replaced
+        expire_on_flush: Whether this property expires on session flush
+        info: Optional data dictionary which will be populated into the MapperProperty.info attribute
+        doc: Optional string that will be applied as the doc on the class-bound descriptor
 
     Returns:
-        A standard ColumnProperty instance with event listeners for fallback handling
+        A MappedSQLExpression instance with event listeners for fallback handling
 
     Example:
         ```python
@@ -46,10 +71,20 @@ def __declare_last__(cls):
                 )
         ```
     """
-    kwargs["deferred"] = deferred
 
-    # Create a standard ColumnProperty
-    prop = column_property(expression, **kwargs)
+    # Create a standard ColumnProperty with all the parameters
+    prop = column_property(
+        expression,
+        *additional_expressions,
+        group=group,
+        deferred=deferred,
+        raiseload=raiseload,
+        comparator_factory=comparator_factory,
+        active_history=active_history,
+        expire_on_flush=expire_on_flush,
+        info=info,
+        doc=doc,
+    )
 
     # Store fallback_value as attribute for later use in event setup
     prop._deferred_fallback_value = fallback_value
@@ -78,6 +113,7 @@ def _setup_fallback_listeners(mapper, key: str, fallback_value: Any):
     @event.listens_for(class_type, "load")
     def _set_deferred_fallback_on_load(target, context):
         """Set fallback value when object is loaded from database"""
+        # Only set fallback if key is not in __dict__ (not loaded)
         if key not in target.__dict__:
             set_committed_value(target, key, fallback_value)
 

tests/test_typing_deferred.py

@@ -0,0 +1,85 @@
+"""
+Test typing compatibility of deferred_column_property.
+
+This test verifies that deferred_column_property has the same typing behavior
+as the original column_property from SQLAlchemy.
+"""
+
+from typing import TYPE_CHECKING, Optional
+
+from sqlalchemy.orm import configure_mappers
+from sqlalchemy.orm import Query, sessionmaker, undefer
+from sqlmodel import Field, SQLModel, create_engine
+
+from sqlmodel.deferred_column import deferred_column_property
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm.attributes import InstrumentedAttribute
+
+
+class User(SQLModel, table=True):
+    id: Optional[int] = Field(default=None, primary_key=True)
+    name: str = Field(default="default")
+    age: Optional[int] = None
+
+    @classmethod
+    def __declare_last__(cls):
+        # Test that deferred_column_property returns proper type
+        cls.computed_age: "InstrumentedAttribute[int]" = deferred_column_property(
+            cls.__table__.c.age * 2, fallback_value=0
+        )
+
+
+def test_typing_compatibility():
+    """Test that typing works correctly with deferred_column_property"""
+    # Force SQLAlchemy to configure mappers
+    configure_mappers()
+
+    # Test that we can use undefer() with deferred_column_property
+    engine = create_engine("sqlite:///:memory:")
+    SQLModel.metadata.create_all(engine)
+
+    SessionLocal = sessionmaker(bind=engine)
+    session = SessionLocal()
+
+    # Create test user
+    user = User(name="Test", age=25)
+    session.add(user)
+    session.commit()
+    user_id = user.id
+
+    # Test loading with undefer (should work without typing errors)
+    query: Query[User] = session.query(User).filter(User.id == user_id)
+
+    loaded_user = query.first()
+    assert loaded_user is not None
+    print(f"Without undefer - loaded_user.computed_age = {loaded_user.computed_age}")
+    print(f"Without undefer - loaded_user.age = {loaded_user.age}")
+
+    # Now try with undefer
+    query_with_undefer: Query[User] = (
+        session.query(User)
+        .options(
+            undefer(User.computed_age)  # This should not cause typing errors
+        )
+        .filter(User.id == user_id)
+    )
+
+    loaded_user_undefer = query_with_undefer.first()
+    assert loaded_user_undefer is not None
+    print(
+        f"With undefer - loaded_user_undefer.computed_age = {loaded_user_undefer.computed_age}"
+    )
+    print(f"With undefer - loaded_user_undefer.age = {loaded_user_undefer.age}")
+
+    # Test fallback behavior when detached
+    session.expunge(loaded_user)  # Detach from session
+    assert loaded_user.computed_age == 0  # Should return fallback value
+
+    session.close()
+
+    print("✅ All typing and functionality tests passed!")
+
+
+if __name__ == "__main__":
+    test_typing_compatibility()

tests/test_typing_final.py

@@ -0,0 +1,78 @@
+"""
+Final typing test for deferred_column_property.
+
+This test verifies that deferred_column_property has the same typing behavior
+as the original column_property from SQLAlchemy.
+"""
+
+from typing import Optional
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import column_property, undefer
+from sqlmodel import Field, SQLModel, Session, select
+
+from sqlmodel.deferred_column import deferred_column_property
+
+
+class TestTyping(SQLModel, table=True):
+    id: Optional[int] = Field(default=None, primary_key=True)
+    salary: int = 50000
+
+    @classmethod
+    def __declare_last__(cls):
+        # Test both standard column_property and our deferred_column_property
+        cls.standard_prop = column_property(cls.__table__.c.salary > 40000)
+        cls.deferred_prop = deferred_column_property(
+            cls.__table__.c.salary > 60000, fallback_value=False
+        )
+
+
+def test_typing_equivalence():
+    """Test that both properties have equivalent typing"""
+
+    # Test that both can be used with undefer() - this verifies typing compatibility
+    engine = create_engine("sqlite:///:memory:")
+    SQLModel.metadata.create_all(engine)
+
+    with Session(engine) as session:
+        user = TestTyping(salary=70000)
+        session.add(user)
+        session.commit()
+        user_id = user.id
+
+    # Test that both properties can be used with undefer (no typing errors)
+    with Session(engine) as session:
+        stmt = (
+            select(TestTyping)
+            .options(
+                undefer(TestTyping.standard_prop),  # Should work without typing errors
+                undefer(TestTyping.deferred_prop),  # Should work without typing errors
+            )
+            .where(TestTyping.id == user_id)
+        )
+
+        loaded_user = session.exec(stmt).one()
+
+        print(f"User salary: {loaded_user.salary}")
+        print(f"Standard prop (salary > 40000): {loaded_user.standard_prop}")
+        print(f"Deferred prop (salary > 60000): {loaded_user.deferred_prop}")
+
+        # Both should return True for salary=70000
+        assert loaded_user.standard_prop is True  # 70000 > 40000
+        assert loaded_user.deferred_prop is True  # 70000 > 60000
+
+        # Test typing compatibility - both should have similar type annotations
+        print(f"Standard prop type: {type(TestTyping.standard_prop)}")
+        print(f"Deferred prop type: {type(TestTyping.deferred_prop)}")
+
+        # Both should have descriptor protocol
+        assert hasattr(TestTyping.standard_prop, "__get__")
+        assert hasattr(TestTyping.deferred_prop, "__get__")
+
+        print("✅ Both properties work correctly with undefer()")
+        print("✅ Both have descriptor protocol")
+        print("✅ Typing compatibility confirmed!")
+
+
+if __name__ == "__main__":
+    test_typing_equivalence()