Merge "KeyFuncDict regression fixes and dataclass fixes" into main

2 files changed, 100 insertions, 35 deletions

diff --git a/lib/sqlalchemy/orm/mapped_collection.py b/lib/sqlalchemy/orm/mapped_collection.py
index a2b085c76..056f14f40 100644
--- a/lib/sqlalchemy/orm/mapped_collection.py
+++ b/lib/sqlalchemy/orm/mapped_collection.py
@@ -7,13 +7,19 @@
 
 from __future__ import annotations
 
+import operator
 from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import Generic
+from typing import List
+from typing import Optional
+from typing import Sequence
+from typing import Tuple
 from typing import Type
 from typing import TYPE_CHECKING
 from typing import TypeVar
+from typing import Union
 
 from . import base
 from .collections import collection
@@ -22,13 +28,9 @@ from .. import util
 from ..sql import coercions
 from ..sql import expression
 from ..sql import roles
+from ..util.typing import Literal
 
 if TYPE_CHECKING:
-    from typing import List
-    from typing import Optional
-    from typing import Sequence
-    from typing import Tuple
-    from typing import Union
 
     from . import AttributeEventToken
     from . import Mapper
@@ -78,7 +80,11 @@ class _PlainColumnGetter(Generic[_KT]):
         if self.composite:
             return tuple(key)
         else:
-            return key[0]
+            obj = key[0]
+            if obj is None:
+                return _UNMAPPED_AMBIGUOUS_NONE
+            else:
+                return obj
 
 
 class _SerializableColumnGetterV2(_PlainColumnGetter[_KT]):
@@ -173,7 +179,7 @@ def column_keyed_dict(
      .. versionadded:: 2.0 an error is raised by default if the attribute
         being used for the dictionary key is determined that it was never
         populated with any value.  The
-        :paramref:`.column_keyed_dict.ignore_unpopulated_attribute`
+        :paramref:`_orm.column_keyed_dict.ignore_unpopulated_attribute`
         parameter may be set which will instead indicate that this condition
         should be ignored, and the append operation silently skipped.
         This is in contrast to the behavior of the 1.x series which would
@@ -193,15 +199,30 @@ def column_keyed_dict(
     )
 
 
+_UNMAPPED_AMBIGUOUS_NONE = object()
+
+
 class _AttrGetter:
-    __slots__ = ("attr_name",)
+    __slots__ = ("attr_name", "getter")
 
     def __init__(self, attr_name: str):
         self.attr_name = attr_name
+        self.getter = operator.attrgetter(attr_name)
 
     def __call__(self, mapped_object: Any) -> Any:
-        dict_ = base.instance_dict(mapped_object)
-        return dict_.get(self.attr_name, base.NO_VALUE)
+        obj = self.getter(mapped_object)
+        if obj is None:
+            state = base.instance_state(mapped_object)
+            mp = state.mapper
+            if self.attr_name in mp.attrs:
+                dict_ = state.dict
+                obj = dict_.get(self.attr_name, base.NO_VALUE)
+                if obj is None:
+                    return _UNMAPPED_AMBIGUOUS_NONE
+            else:
+                return _UNMAPPED_AMBIGUOUS_NONE
+
+        return obj
 
     def __reduce__(self) -> Tuple[Type[_AttrGetter], Tuple[str]]:
         return _AttrGetter, (self.attr_name,)
@@ -240,7 +261,7 @@ def attribute_keyed_dict(
      .. versionadded:: 2.0 an error is raised by default if the attribute
         being used for the dictionary key is determined that it was never
         populated with any value.  The
-        :paramref:`.attribute_keyed_dict.ignore_unpopulated_attribute`
+        :paramref:`_orm.attribute_keyed_dict.ignore_unpopulated_attribute`
         parameter may be set which will instead indicate that this condition
         should be ignored, and the append operation silently skipped.
         This is in contrast to the behavior of the 1.x series which would
@@ -291,7 +312,7 @@ def keyfunc_mapping(
         being used for the dictionary key returns
         :attr:`.LoaderCallableStatus.NO_VALUE`, which in an ORM attribute
         context indicates an attribute that was never populated with any value.
-        The :paramref:`.mapped_collection.ignore_unpopulated_attribute`
+        The :paramref:`_orm.mapped_collection.ignore_unpopulated_attribute`
         parameter may be set which will instead indicate that this condition
         should be ignored, and the append operation silently skipped. This is
         in contrast to the behavior of the 1.x series which would erroneously
@@ -334,7 +355,7 @@ class KeyFuncDict(Dict[_KT, _VT]):
     def __init__(
         self,
         keyfunc: _F,
-        *,
+        *dict_args: Any,
         ignore_unpopulated_attribute: bool = False,
     ) -> None:
         """Create a new collection with keying provided by keyfunc.
@@ -352,6 +373,7 @@ class KeyFuncDict(Dict[_KT, _VT]):
         """
         self.keyfunc = keyfunc
         self.ignore_unpopulated_attribute = ignore_unpopulated_attribute
+        super().__init__(*dict_args)
 
     @classmethod
     def _unreduce(
@@ -369,36 +391,56 @@ class KeyFuncDict(Dict[_KT, _VT]):
     ]:
         return (KeyFuncDict._unreduce, (self.keyfunc, dict(self)))
 
+    @util.preload_module("sqlalchemy.orm.attributes")
     def _raise_for_unpopulated(
-        self, value: _KT, initiator: Optional[AttributeEventToken]
+        self,
+        value: _KT,
+        initiator: Union[AttributeEventToken, Literal[None, False]] = None,
+        *,
+        warn_only: bool,
     ) -> None:
         mapper = base.instance_state(value).mapper
 
-        if initiator is None:
+        attributes = util.preloaded.orm_attributes
+
+        if not isinstance(initiator, attributes.AttributeEventToken):
             relationship = "unknown relationship"
-        else:
+        elif initiator.key in mapper.attrs:
             relationship = f"{mapper.attrs[initiator.key]}"
-
-        raise sa_exc.InvalidRequestError(
-            f"In event triggered from population of attribute {relationship} "
-            "(likely from a backref), "
-            f"can't populate value in KeyFuncDict; "
-            "dictionary key "
-            f"derived from {base.instance_str(value)} is not "
-            f"populated. Ensure appropriate state is set up on "
-            f"the {base.instance_str(value)} object "
-            f"before assigning to the {relationship} attribute. "
-            f"To skip this assignment entirely, "
-            f'Set the "ignore_unpopulated_attribute=True" '
-            f"parameter on the mapped collection factory."
-        )
+        else:
+            relationship = initiator.key
+
+        if warn_only:
+            util.warn(
+                f"Attribute keyed dictionary value for "
+                f"attribute '{relationship}' was None; this will raise "
+                "in a future release. "
+                f"To skip this assignment entirely, "
+                f'Set the "ignore_unpopulated_attribute=True" '
+                f"parameter on the mapped collection factory."
+            )
+        else:
+            raise sa_exc.InvalidRequestError(
+                "In event triggered from population of "
+                f"attribute '{relationship}' "
+                "(potentially from a backref), "
+                f"can't populate value in KeyFuncDict; "
+                "dictionary key "
+                f"derived from {base.instance_str(value)} is not "
+                f"populated. Ensure appropriate state is set up on "
+                f"the {base.instance_str(value)} object "
+                f"before assigning to the {relationship} attribute. "
+                f"To skip this assignment entirely, "
+                f'Set the "ignore_unpopulated_attribute=True" '
+                f"parameter on the mapped collection factory."
+            )
 
     @collection.appender  # type: ignore[misc]
     @collection.internally_instrumented  # type: ignore[misc]
     def set(
         self,
         value: _KT,
-        _sa_initiator: Optional[AttributeEventToken] = None,
+        _sa_initiator: Union[AttributeEventToken, Literal[None, False]] = None,
     ) -> None:
         """Add an item by value, consulting the keyfunc for the key."""
 
@@ -406,7 +448,17 @@ class KeyFuncDict(Dict[_KT, _VT]):
 
         if key is base.NO_VALUE:
             if not self.ignore_unpopulated_attribute:
-                self._raise_for_unpopulated(value, _sa_initiator)
+                self._raise_for_unpopulated(
+                    value, _sa_initiator, warn_only=False
+                )
+            else:
+                return
+        elif key is _UNMAPPED_AMBIGUOUS_NONE:
+            if not self.ignore_unpopulated_attribute:
+                self._raise_for_unpopulated(
+                    value, _sa_initiator, warn_only=True
+                )
+                key = None
             else:
                 return
 
@@ -417,7 +469,7 @@ class KeyFuncDict(Dict[_KT, _VT]):
     def remove(
         self,
         value: _KT,
-        _sa_initiator: Optional[AttributeEventToken] = None,
+        _sa_initiator: Union[AttributeEventToken, Literal[None, False]] = None,
     ) -> None:
         """Remove an item by value, consulting the keyfunc for the key."""
 
@@ -425,8 +477,18 @@ class KeyFuncDict(Dict[_KT, _VT]):
 
         if key is base.NO_VALUE:
             if not self.ignore_unpopulated_attribute:
-                self._raise_for_unpopulated(value, _sa_initiator)
+                self._raise_for_unpopulated(
+                    value, _sa_initiator, warn_only=False
+                )
             return
+        elif key is _UNMAPPED_AMBIGUOUS_NONE:
+            if not self.ignore_unpopulated_attribute:
+                self._raise_for_unpopulated(
+                    value, _sa_initiator, warn_only=True
+                )
+                key = None
+            else:
+                return
 
         # Let self[key] raise if key is not in this collection
         # testlib.pragma exempt:__ne__
@@ -444,9 +506,10 @@ def _mapped_collection_cls(
     keyfunc: _F, ignore_unpopulated_attribute: bool
 ) -> Type[KeyFuncDict[_KT, _KT]]:
     class _MKeyfuncMapped(KeyFuncDict[_KT, _KT]):
-        def __init__(self) -> None:
+        def __init__(self, *dict_args: Any) -> None:
             super().__init__(
                 keyfunc,
+                *dict_args,
                 ignore_unpopulated_attribute=ignore_unpopulated_attribute,
             )
 
diff --git a/lib/sqlalchemy/util/preloaded.py b/lib/sqlalchemy/util/preloaded.py
index 666a208c0..f3609c8e4 100644
--- a/lib/sqlalchemy/util/preloaded.py
+++ b/lib/sqlalchemy/util/preloaded.py
@@ -29,6 +29,7 @@ if TYPE_CHECKING:
     from sqlalchemy.engine import reflection as _engine_reflection
     from sqlalchemy.engine import result as _engine_result
     from sqlalchemy.engine import url as _engine_url
+    from sqlalchemy.orm import attributes as _orm_attributes
     from sqlalchemy.orm import base as _orm_base
     from sqlalchemy.orm import clsregistry as _orm_clsregistry
     from sqlalchemy.orm import decl_api as _orm_decl_api
@@ -65,6 +66,7 @@ if TYPE_CHECKING:
     orm_clsregistry = _orm_clsregistry
     orm_base = _orm_base
     orm = _orm
+    orm_attributes = _orm_attributes
     orm_decl_api = _orm_decl_api
     orm_decl_base = _orm_decl_base
     orm_descriptor_props = _orm_descriptor_props