From f560250b2ac278ccd8513a0b0b1511c59fea4af8 Mon Sep 17 00:00:00 2001
From: Federico Caselli <cfederico87@gmail.com>
Date: Wed, 1 Sep 2021 22:23:19 +0200
Subject: [PATCH] Add __get__ to the column and relationship to improve usage
 when the plugin is not supported, for example when using other type checkers.

---
 sqlalchemy-stubs/orm/relationships.pyi |  8 +++++
 sqlalchemy-stubs/sql/schema.pyi        |  7 +++++
 test/files/orm_columns.py              | 42 ++++++++++++++++++++++++++
 3 files changed, 57 insertions(+)
 create mode 100644 test/files/orm_columns.py

diff --git a/sqlalchemy-stubs/orm/relationships.pyi b/sqlalchemy-stubs/orm/relationships.pyi
index 925dd9d..d15450a 100644
--- a/sqlalchemy-stubs/orm/relationships.pyi
+++ b/sqlalchemy-stubs/orm/relationships.pyi
@@ -3,6 +3,7 @@ from typing import Any
 from typing import Callable
 from typing import MutableMapping
 from typing import Optional
+from typing import overload
 from typing import Sequence
 from typing import Tuple
 from typing import Type
@@ -13,6 +14,7 @@ from typing_extensions import Literal
 
 from . import _BackrefResult
 from . import attributes as attributes
+from .attributes import Mapped
 from .base import state_str as state_str
 from .interfaces import MANYTOMANY as MANYTOMANY
 from .interfaces import MANYTOONE as MANYTOONE
@@ -188,6 +190,12 @@ class RelationshipProperty(StrategizedProperty[_T]):
     def cascade(self) -> CascadeOptions: ...
     @cascade.setter
     def cascade(self, cascade: Sequence[str]) -> None: ...
+    # NOTE: his doesn't exist at runtime, it's used when not using the plugin
+    # makes the inferred type of class attrs Mapped, Any for instance ones
+    @overload
+    def __get__(self, instance: None, owner: Any) -> Mapped: ...
+    @overload
+    def __get__(self, instance: object, owner: Any) -> Any: ...
 
 class JoinCondition:
     parent_persist_selectable: Any = ...
diff --git a/sqlalchemy-stubs/sql/schema.pyi b/sqlalchemy-stubs/sql/schema.pyi
index 244dd8f..1616699 100644
--- a/sqlalchemy-stubs/sql/schema.pyi
+++ b/sqlalchemy-stubs/sql/schema.pyi
@@ -34,6 +34,7 @@ from .selectable import TableClause
 from .. import util
 from ..engine import Connection
 from ..engine import Engine
+from ..orm import Mapped
 from ..util import langhelpers
 
 RETAIN_SCHEMA: langhelpers._symbol
@@ -246,6 +247,12 @@ class Column(DialectKWArgs, SchemaItem, ColumnClause[_TE]):
     def references(self, column: Column[Any]) -> bool: ...
     def append_foreign_key(self, fk: ForeignKey) -> None: ...
     def copy(self: _CO, **kw: Any) -> _CO: ...
+    # NOTE: his doesn't exist at runtime, it's used when not using the plugin
+    # makes the inferred type of orm class attrs Mapped, Any for instance ones
+    @overload
+    def __get__(self, instance: None, owner: Any) -> Mapped: ...
+    @overload
+    def __get__(self, instance: object, owner: Any) -> Any: ...
 
 class ForeignKey(DialectKWArgs, SchemaItem):
     __visit_name__: str = ...
diff --git a/test/files/orm_columns.py b/test/files/orm_columns.py
new file mode 100644
index 0000000..ea859b1
--- /dev/null
+++ b/test/files/orm_columns.py
@@ -0,0 +1,42 @@
+from typing import List
+from typing import Optional
+
+from sqlalchemy import Column
+from sqlalchemy import Integer
+from sqlalchemy import ForeignKey
+from sqlalchemy import Text
+from sqlalchemy.orm import registry
+from sqlalchemy.orm import relationship
+from sqlalchemy.orm import Mapped
+
+reg: registry = registry()
+
+
+@reg.mapped
+class Model:
+    __tablename__ = "model"
+
+    id = Column(Integer, primary_key=True)
+    name = Column(Text)
+
+
+@reg.mapped
+class Other:
+    __tablename__ = "other"
+    id = Column(Integer, primary_key=True)
+    model_id = Column(Integer, ForeignKey("model.id"))
+
+    model = relationship(Model, uselist=False)
+
+
+col_id: Column[Integer] = Model.id
+col_name: Column[Text] = Model.name
+other_id: Column[Integer] = Other.id
+other_mid: Column[Integer] = Other.model_id
+other_model: Mapped[Model] = Other.model
+
+inst_id: Optional[int] = Model().id
+inst_name: Optional[str] = Model().name
+inst_other_id: Optional[int] = Other().id
+inst_other_mid: Optional[int] = Other().model_id
+inst_other_model: Optional[Model] = Other().model