Merge "Improve support for enum in mapped classes" into main

2 files changed, 83 insertions, 0 deletions

diff --git a/lib/sqlalchemy/sql/sqltypes.py b/lib/sqlalchemy/sql/sqltypes.py
index 92708bb40..1c58bea20 100644
--- a/lib/sqlalchemy/sql/sqltypes.py
+++ b/lib/sqlalchemy/sql/sqltypes.py
@@ -59,6 +59,7 @@ from .. import util
 from ..engine import processors
 from ..util import langhelpers
 from ..util import OrderedDict
+from ..util.typing import GenericProtocol
 from ..util.typing import Literal
 
 if TYPE_CHECKING:
@@ -1489,6 +1490,28 @@ class Enum(String, SchemaType, Emulated, TypeEngine[Union[str, enum.Enum]]):
             self.enum_class = None
             return enums, enums
 
+    def _resolve_for_literal(self, value: Any) -> Enum:
+        typ = self._resolve_for_python_type(type(value), type(value))
+        assert typ is not None
+        return typ
+
+    def _resolve_for_python_type(
+        self,
+        python_type: Type[Any],
+        matched_on: Union[GenericProtocol[Any], Type[Any]],
+    ) -> Optional[Enum]:
+        if not issubclass(python_type, enum.Enum):
+            return None
+        return cast(
+            Enum,
+            util.constructor_copy(
+                self,
+                self._generic_type_affinity,
+                python_type,
+                length=NO_ARG if self.length == 0 else self.length,
+            ),
+        )
+
     def _setup_for_values(self, values, objects, kw):
         self.enums = list(values)
 
@@ -3674,6 +3697,7 @@ _type_map: Dict[Type[Any], TypeEngine[Any]] = {
     type(None): NULLTYPE,
     bytes: LargeBinary(),
     str: _STRING,
+    enum.Enum: Enum(enum.Enum),
 }
 
 
diff --git a/lib/sqlalchemy/sql/type_api.py b/lib/sqlalchemy/sql/type_api.py
index c3768c6c6..b395e6796 100644
--- a/lib/sqlalchemy/sql/type_api.py
+++ b/lib/sqlalchemy/sql/type_api.py
@@ -35,6 +35,7 @@ from .operators import ColumnOperators
 from .visitors import Visitable
 from .. import exc
 from .. import util
+from ..util.typing import flatten_generic
 from ..util.typing import Protocol
 from ..util.typing import TypedDict
 from ..util.typing import TypeGuard
@@ -55,6 +56,7 @@ if typing.TYPE_CHECKING:
     from .sqltypes import STRINGTYPE as STRINGTYPE  # noqa: F401
     from .sqltypes import TABLEVALUE as TABLEVALUE  # noqa: F401
     from ..engine.interfaces import Dialect
+    from ..util.typing import GenericProtocol
 
 _T = TypeVar("_T", bound=Any)
 _T_co = TypeVar("_T_co", bound=Any, covariant=True)
@@ -712,9 +714,66 @@ class TypeEngine(Visitable, Generic[_T]):
 
         .. versionadded:: 1.4.30 or 2.0
 
+        TODO: this should be part of public API
+
+        .. seealso::
+
+            :meth:`.TypeEngine._resolve_for_python_type`
+
         """
         return self
 
+    def _resolve_for_python_type(
+        self: SelfTypeEngine,
+        python_type: Type[Any],
+        matched_on: Union[GenericProtocol[Any], Type[Any]],
+    ) -> Optional[SelfTypeEngine]:
+        """given a Python type (e.g. ``int``, ``str``, etc. ) return an
+        instance of this :class:`.TypeEngine` that's appropriate for this type.
+
+        An additional argument ``matched_on`` is passed, which indicates an
+        entry from the ``__mro__`` of the given ``python_type`` that more
+        specifically matches how the caller located this :class:`.TypeEngine`
+        object.   Such as, if a lookup of some kind links the ``int`` Python
+        type to the :class:`.Integer` SQL type, and the original object
+        was some custom subclass of ``int`` such as ``MyInt(int)``, the
+        arguments passed would be ``(MyInt, int)``.
+
+        If the given Python type does not correspond to this
+        :class:`.TypeEngine`, or the Python type is otherwise ambiguous, the
+        method should return None.
+
+        For simple cases, the method checks that the ``python_type``
+        and ``matched_on`` types are the same (i.e. not a subclass), and
+        returns self; for all other cases, it returns ``None``.
+
+        The initial use case here is for the ORM to link user-defined
+        Python standard library ``enum.Enum`` classes to the SQLAlchemy
+        :class:`.Enum` SQL type when constructing ORM Declarative mappings.
+
+        :param python_type: the Python type we want to use
+        :param matched_on: the Python type that led us to choose this
+         particular :class:`.TypeEngine` class, which would be a supertype
+         of ``python_type``.   By default, the request is rejected if
+         ``python_type`` doesn't match ``matched_on`` (None is returned).
+
+        .. versionadded:: 2.0.0b4
+
+        TODO: this should be part of public API
+
+        .. seealso::
+
+            :meth:`.TypeEngine._resolve_for_literal`
+
+        """
+
+        matched_on = flatten_generic(matched_on)
+
+        if python_type is not matched_on:
+            return None
+
+        return self
+
     @util.ro_memoized_property
     def _type_affinity(self) -> Optional[Type[TypeEngine[_T]]]:
         """Return a rudimental 'affinity' value expressing the general class