From: Bryan Forbes <bryan@reigndropsfall.net>
Date: Thu, 8 Apr 2021 16:27:17 +0000 (-0400)
Subject: Support `TypeDecorator` subclasses in `Column()` declarations
X-Git-Tag: rel_1_4_7~8^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=ce6f8ec5f5b76267e75b63827142ceda2276d918;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

Support `TypeDecorator` subclasses in `Column()` declarations

<!-- Provide a general summary of your proposed changes in the Title field above -->

### Description
Currently, the plugin resolves `TypeDecorator` subclasses in `Column()` declarations to `Mapping[_T]` instead of the correct type in the class declaration.

### Checklist
<!-- go over following points. check them with an `x` if they do apply, (they turn into clickable checkboxes once the PR is submitted, so no need to do everything at once)

-->

This pull request is:

- [ ] A documentation / typographical error fix
	- Good to go, no issue or tests are needed
- [X] A short code fix
	- please include the issue number, and create an issue if none exists, which
	  must include a complete example of the issue.  one line code fixes without an
	  issue and demonstration will not be accepted.
	- Please include: `Fixes: #<issue number>` in the commit message
	- please include tests.   one line code fixes without tests will not be accepted.
- [ ] A new feature implementation
	- please include the issue number, and create an issue if none exists, which must
	  include a complete example of how the feature would look.
	- Please include: `Fixes: #<issue number>` in the commit message
	- please include tests.

**Have a nice day!**

Closes: #6223
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/6223
Pull-request-sha: 5136bc4b6333e868cc47f1b2dcc58716a40cadca

Change-Id: Idc865bf7320f8ea3054c28dea095b693fe112753
---

diff --git a/doc/build/changelog/unreleased_14/mypy_typedec.rst b/doc/build/changelog/unreleased_14/mypy_typedec.rst
new file mode 100644
index 0000000000..f4f7f3c195
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/mypy_typedec.rst
@@ -0,0 +1,7 @@
+.. change::
+    :tags: bug, mypy
+
+    Fixed issue in Mypy plugin where the plugin wasn’t inferring the correct 
+    type for columns of subclasses that don’t directly descend from 
+    ``TypeEngine``, in particular that of  ``TypeDecorator`` and 
+    ``UserDefinedType``.
diff --git a/lib/sqlalchemy/ext/mypy/infer.py b/lib/sqlalchemy/ext/mypy/infer.py
index 49dd9fb743..f0f6be36f3 100644
--- a/lib/sqlalchemy/ext/mypy/infer.py
+++ b/lib/sqlalchemy/ext/mypy/infer.py
@@ -10,6 +10,7 @@ from typing import Union
 
 from mypy import nodes
 from mypy import types
+from mypy.maptype import map_instance_to_supertype
 from mypy.messages import format_type
 from mypy.nodes import AssignmentStmt
 from mypy.nodes import CallExpr
@@ -413,9 +414,11 @@ def _extract_python_type_from_typeengine(
             n = api.lookup_fully_qualified("builtins.str")
             return Instance(n.node, [])
 
-    for mr in node.mro:
-        if mr.bases:
-            for base_ in mr.bases:
-                if base_.type.fullname == "sqlalchemy.sql.type_api.TypeEngine":
-                    return base_.args[-1]
-    assert False, "could not extract Python type from node: %s" % node
+    assert node.has_base("sqlalchemy.sql.type_api.TypeEngine"), (
+        "could not extract Python type from node: %s" % node
+    )
+    type_engine = map_instance_to_supertype(
+        Instance(node, []),
+        api.modules["sqlalchemy.sql.type_api"].names["TypeEngine"].node,
+    )
+    return type_engine.args[-1]
diff --git a/test/ext/mypy/files/type_decorator.py b/test/ext/mypy/files/type_decorator.py
new file mode 100644
index 0000000000..83f603c63e
--- /dev/null
+++ b/test/ext/mypy/files/type_decorator.py
@@ -0,0 +1,45 @@
+from typing import Any
+from typing import Optional
+
+from sqlalchemy import Column
+from sqlalchemy import Integer
+from sqlalchemy import String
+from sqlalchemy import TypeDecorator
+from sqlalchemy.ext.declarative import declarative_base
+
+Base = declarative_base()
+
+
+class IntToStr(TypeDecorator[int]):
+    impl = String
+
+    def process_bind_param(
+        self,
+        value: Any,
+        dialect: Any,
+    ) -> Optional[str]:
+        return str(value) if value is not None else value
+
+    def process_result_value(
+        self,
+        value: Any,
+        dialect: Any,
+    ) -> Optional[int]:
+        return int(value) if value is not None else value
+
+    def copy(self, **kwargs: Any) -> "IntToStr":
+        return IntToStr(self.impl.length)
+
+
+class Thing(Base):
+    __tablename__ = "things"
+
+    id: int = Column(Integer, primary_key=True)
+    intToStr: int = Column(IntToStr)
+
+
+t1 = Thing(intToStr=5)
+
+i5: int = t1.intToStr
+
+t1.intToStr = 8