Skip to content

Commit b0b74dc

Browse files
CaselITCaselIT
committed
Improve typing story for core from clauses.
Most :class:`_sql.FromClause` subclasses are not generic on :class:`.TypedColumns` subclasses, that can be used to type their :attr:`_sql.FromClause.c` collection. This applied to :class:`_schema.Table`, :class:`_sql.Join`, :class:`_sql.Subquery`, :class:`_sql.CTE` and more. Fixes: #13085 Change-Id: I724aca887a85c4a401df875903eda12125066680
1 parent d7ade42 commit b0b74dc

25 files changed

Lines changed: 2253 additions & 119 deletions

‎doc/build/changelog/migration_21.rst‎

Lines changed: 87 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ dataclass-level default (i.e. set using any of the
198198
:paramref:`_orm.column_property.default`, or :paramref:`_orm.deferred.default`
199199
parameters) is directed to be delivered at the
200200
Python :term:`descriptor` level using mechanisms in SQLAlchemy's attribute
201-
system that normally return ``None`` for un-popualted columns, so that even though the default is not
201+
system that normally return ``None`` for un-populated columns, so that even though the default is not
202202
populated into ``__dict__``, it's still delivered when the attribute is
203203
accessed. This behavior is based on what Python dataclasses itself does
204204
when a default is indicated for a field that also includes ``init=False``.
@@ -721,6 +721,92 @@ up front, which would be verbose and not automatic.
721721
722722
:ticket:`10635`
723723

724+
.. _change_13085:
725+
726+
Better type checker integration for Core froms, like Table
727+
----------------------------------------------------------
728+
729+
SQLAlchemy 2.1 changes :class:`_schema.Table`, along with most
730+
:class:`_sql.FromClause` subclasses, to be generic on the column collection,
731+
providing the option for better static type checking support.
732+
By declaring the columns using a :class:`_schema.TypedColumns` subclass and
733+
providing it to the :class:`_schema.Table` instance, IDEs and type checkers
734+
can infer the exact types of columns when accessing them via the
735+
:attr:`_schema.Table.c` attribute, enabling better autocomplete and type validation.
736+
737+
Example usage::
738+
739+
from sqlalchemy import Table, TypedColumns, Column, Integer, MetaData, select
740+
741+
742+
class user_cols(TypedColumns):
743+
id = Column(Integer, primary_key=True)
744+
name: Column[str]
745+
age: Column[int]
746+
747+
# optional, used to infer the select types when selecting the table
748+
__row_pos__: tuple[int, str, int]
749+
750+
751+
metadata = MetaData()
752+
user = Table("user", metadata, user_cols)
753+
754+
# Type checkers now understand the column types when selecting single columns
755+
stmt = select(user.c.id, user.c.name) # Inferred as Select[int, str]
756+
757+
# and also when selecting the whole table, when __row_pos__ is present
758+
stmt = select(user) # Inferred as Select[int, str, int]
759+
760+
The optional :attr:`sqlalchemy.sql._annotated_cols.HasRowPos.__row_pos__` annotation
761+
is used to infer the types of a select when selecting the table directly.
762+
763+
Columns can be declared in :class:`.TypedColumns` subclasses by instantiating
764+
them directly or by using only a type annotations, that will be inferred when
765+
generating a :class:`_schema.Table`.
766+
767+
Other :class:`_sql.FromClause`, like :class:`_sql.Join`, :class:`_sql.CTE`, etc, can be made
768+
generic using the :meth:`_sql.FromClause.with_cols` method::
769+
770+
# using with_cols the ``c`` collection of the cte has typed tables
771+
cte = user.select().cte().with_cols(user_cols)
772+
773+
ORM Integration
774+
^^^^^^^^^^^^^^^
775+
776+
This functionality also offers some integration with the ORM, by using
777+
:class:`_orm.MappedColumn` annotated attributes in the ORM model and
778+
:func:`_orm.as_typed_table` to get an annotated :class:`_sql.FromClause`::
779+
780+
from sqlalchemy import TypedColumns
781+
from sqlalchemy.orm import DeclarativeBase, mapped_column
782+
from sqlalchemy.orm import MappedColumn, as_typed_table
783+
784+
785+
class Base(DeclarativeBase):
786+
pass
787+
788+
789+
class A(Base):
790+
__tablename__ = "a"
791+
__typed_cols__: "a_cols"
792+
793+
id: MappedColumn[int] = mapped_column(primary_key=True)
794+
data: MappedColumn[str]
795+
796+
797+
class a_cols(A, TypedColumns):
798+
pass
799+
800+
801+
# table_a is annotated as FromClause[a_cols], and is just A.__table__
802+
table_a = as_typed_table(A)
803+
804+
For proper typing integration :class:`_orm.MappedColumn` should be used
805+
to annotate the single columns, since it's a more specific annotation than
806+
the usual :class:`_orm.Mapped` used for ORM attributes.
807+
808+
:ticket:`13085`
809+
724810
.. _change_8601:
725811

726812
``filter_by()`` now searches across all FROM clause entities

‎doc/build/changelog/unreleased_21/13085.rst‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
.. change::
2+
:tags: schema, usecase
3+
:tickets: 13085
4+
5+
Most :class:`_sql.FromClause` subclasses are now generic on
6+
:class:`_schema.TypedColumns` subclasses, that can be used to type their
7+
:attr:`_sql.FromClause.c` collection.
8+
This applied to :class:`_schema.Table`, :class:`_sql.Join`,
9+
:class:`_sql.Subquery`, :class:`_sql.CTE` and more.
10+
11+
.. seealso::
12+
13+
:ref:`change_13085`

‎doc/build/core/metadata.rst‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -916,3 +916,13 @@ Column, Table, MetaData API
916916
.. autoclass:: Table
917917
:members:
918918
:inherited-members:
919+
920+
.. autoclass:: TypedColumns
921+
:members:
922+
923+
.. autoclass:: Named
924+
:members:
925+
926+
.. autoclass:: sqlalchemy.sql._annotated_cols.HasRowPos
927+
:special-members: __row_pos__
928+

‎doc/build/orm/mapping_api.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,4 +150,6 @@ Class Mapping API
150150

151151
.. autofunction:: unmapped_dataclass
152152

153+
.. autofunction:: as_typed_table
154+
153155

‎doc/build/tutorial/metadata.rst‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,39 @@ parameter.
197197
related column, in the above example the :class:`_types.Integer` datatype
198198
of the ``user_account.id`` column.
199199

200-
In the next section we will emit the completed DDL for the ``user`` and
200+
Using :class:`.TypedColumns` to get a better typing experience
201+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
202+
203+
A SQLAlchemy :class:`_schema.Table` can also be defined using a
204+
:class:`_schema.TypedColumns` to offers better integration with type checker and IDEs.
205+
The tables defined above could be declared as follows::
206+
207+
>>> from sqlalchemy import Named, TypedColumns, Table
208+
>>> other_meta = MetaData()
209+
>>> class user_cols(TypedColumns):
210+
... id: Named[int] = Column(primary_key=True)
211+
... name: Named[str | None] = Column(String(30))
212+
... fullname: Named[str | None]
213+
214+
>>> typed_user_table = Table("user_account", other_meta, user_cols)
215+
216+
>>> class address_cols(TypedColumns):
217+
... id: Named[int] = Column(primary_key=True)
218+
... user_id: Named[int] = Column(ForeignKey("user_account.id"))
219+
... email_address: Named[str]
220+
... __row_pos__: tuple[int, int, str]
221+
222+
>>> typed_address_table = Table("address", other_meta, address_cols)
223+
224+
The columns are defined by subclassing :class:`.TypedColumns`, so that
225+
static type checkers can understand what columns are present in the
226+
:attr:`_schema.Table.c` collection. Functionally the two methods of defining
227+
the metadata objects are equivalent.
228+
The optional ``__row_pos__`` annotation is an aid to type checker so that
229+
they can correctly suggest the type to apply when selecting from the complete
230+
table, without specifying the single columns.
231+
232+
In the next section we will emit the completed DDL for the ``user_account`` and
201233
``address`` table to see the completed result.
202234

203235
.. _tutorial_emitting_ddl:
@@ -576,7 +608,7 @@ are found to be present already:
576608
.. _tutorial_table_reflection:
577609

578610
Table Reflection
579-
-------------------------------
611+
----------------
580612

581613
.. topic:: Optional Section
582614

‎lib/sqlalchemy/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,9 +78,11 @@
7878
from .schema import Index as Index
7979
from .schema import insert_sentinel as insert_sentinel
8080
from .schema import MetaData as MetaData
81+
from .schema import Named as Named
8182
from .schema import PrimaryKeyConstraint as PrimaryKeyConstraint
8283
from .schema import Sequence as Sequence
8384
from .schema import Table as Table
85+
from .schema import TypedColumns as TypedColumns
8486
from .schema import UniqueConstraint as UniqueConstraint
8587
from .sql import ColumnExpressionArgument as ColumnExpressionArgument
8688
from .sql import NotNullable as NotNullable

‎lib/sqlalchemy/orm/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
from .context import QueryContext as QueryContext
5757
from .decl_api import add_mapped_attribute as add_mapped_attribute
5858
from .decl_api import as_declarative as as_declarative
59+
from .decl_api import as_typed_table as as_typed_table
5960
from .decl_api import declarative_base as declarative_base
6061
from .decl_api import declarative_mixin as declarative_mixin
6162
from .decl_api import DeclarativeBase as DeclarativeBase

‎lib/sqlalchemy/orm/base.py‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,14 @@
2727
from typing import Union
2828

2929
from . import exc
30+
from ._typing import _O
3031
from ._typing import insp_is_mapper
3132
from .. import exc as sa_exc
3233
from .. import inspection
3334
from .. import util
3435
from ..sql import roles
36+
from ..sql._typing import _T
37+
from ..sql._typing import _T_co
3538
from ..sql.elements import SQLColumnExpression
3639
from ..sql.elements import SQLCoreOperations
3740
from ..util import FastIntFlag
@@ -46,18 +49,16 @@
4649
from .instrumentation import ClassManager
4750
from .interfaces import PropComparator
4851
from .mapper import Mapper
52+
from .properties import MappedColumn
4953
from .state import InstanceState
5054
from .util import AliasedClass
5155
from .writeonly import WriteOnlyCollection
56+
from ..sql._annotated_cols import TypedColumns
5257
from ..sql._typing import _ColumnExpressionArgument
5358
from ..sql._typing import _InfoType
5459
from ..sql.elements import ColumnElement
5560
from ..sql.operators import OperatorType
56-
57-
_T = TypeVar("_T", bound=Any)
58-
_T_co = TypeVar("_T_co", bound=Any, covariant=True)
59-
60-
_O = TypeVar("_O", bound=object)
61+
from ..sql.schema import Column
6162

6263

6364
class LoaderCallableStatus(Enum):
@@ -804,6 +805,11 @@ class Mapped(
804805

805806
if typing.TYPE_CHECKING:
806807

808+
@overload
809+
def __get__( # type: ignore[misc]
810+
self: MappedColumn[_T_co], instance: TypedColumns, owner: Any
811+
) -> Column[_T_co]: ...
812+
807813
@overload
808814
def __get__(
809815
self, instance: None, owner: Any
@@ -814,7 +820,7 @@ def __get__(self, instance: object, owner: Any) -> _T_co: ...
814820

815821
def __get__(
816822
self, instance: Optional[object], owner: Any
817-
) -> Union[InstrumentedAttribute[_T_co], _T_co]: ...
823+
) -> Union[InstrumentedAttribute[_T_co], Column[_T_co], _T_co]: ...
818824

819825
@classmethod
820826
def _empty_constructor(cls, arg1: Any) -> Mapped[_T_co]: ...

‎lib/sqlalchemy/orm/decl_api.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from typing import Mapping
2424
from typing import Optional
2525
from typing import overload
26+
from typing import Protocol
2627
from typing import Set
2728
from typing import Tuple
2829
from typing import Type
@@ -52,6 +53,7 @@
5253
from .decl_base import _DeferredDeclarativeConfig
5354
from .decl_base import _del_attribute
5455
from .decl_base import _ORMClassConfigurator
56+
from .decl_base import MappedClassProtocol
5557
from .descriptor_props import Composite
5658
from .descriptor_props import Synonym
5759
from .descriptor_props import Synonym as _orm_synonym
@@ -65,6 +67,7 @@
6567
from ..event import dispatcher
6668
from ..event import EventTarget
6769
from ..sql import sqltypes
70+
from ..sql._annotated_cols import _TC
6871
from ..sql.base import _NoArg
6972
from ..sql.elements import SQLCoreOperations
7073
from ..sql.schema import MetaData
@@ -753,6 +756,100 @@ def _sa_inspect_instance(self) -> InstanceState[Self]: ...
753756
def __init__(self, **kw: Any): ...
754757

755758

759+
class MappedClassWithTypedColumnsProtocol(Protocol[_TC]):
760+
"""An ORM mapped class that also defines in the ``__typed_cols__``
761+
attribute its .
762+
"""
763+
764+
__typed_cols__: _TC
765+
"""The :class:`_schema.TypedColumns` of this ORM mapped class."""
766+
767+
__name__: ClassVar[str]
768+
__mapper__: ClassVar[Mapper[Any]]
769+
__table__: ClassVar[FromClause]
770+
771+
772+
@overload
773+
def as_typed_table(
774+
cls: type[MappedClassWithTypedColumnsProtocol[_TC]], /
775+
) -> FromClause[_TC]: ...
776+
777+
778+
@overload
779+
def as_typed_table(
780+
cls: MappedClassProtocol[Any], typed_columns_cls: type[_TC], /
781+
) -> FromClause[_TC]: ...
782+
783+
784+
def as_typed_table(
785+
cls: (
786+
MappedClassProtocol[Any]
787+
| type[MappedClassWithTypedColumnsProtocol[Any]]
788+
),
789+
typed_columns_cls: Any = None,
790+
/,
791+
) -> FromClause[Any]:
792+
"""Return a typed :class:`_sql.FromClause` from the give ORM model.
793+
794+
This function is just a typing help, at runtime it just returns the
795+
``__table__`` attribute of the provided ORM model.
796+
797+
It's usually called providing both the ORM model and the
798+
:class:`_schema.TypedColumns` class. Single argument calls are supported
799+
if the ORM model class provides an annotation pointing to its
800+
:class:`_schema.TypedColumns` in the ``__typed_cols__`` attribute.
801+
802+
803+
Example usage::
804+
805+
from sqlalchemy import TypedColumns
806+
from sqlalchemy.orm import DeclarativeBase, mapped_column
807+
from sqlalchemy.orm import MappedColumn, as_typed_table
808+
809+
810+
class Base(DeclarativeBase):
811+
pass
812+
813+
814+
class A(Base):
815+
__tablename__ = "a"
816+
817+
id: MappedColumn[int] = mapped_column(primary_key=True)
818+
data: MappedColumn[str]
819+
820+
821+
class a_cols(A, TypedColumns):
822+
pass
823+
824+
825+
# table_a is annotated as FromClause[a_cols]
826+
table_a = as_typed_table(A, a_cols)
827+
828+
829+
class B(Base):
830+
__tablename__ = "b"
831+
__typed_cols__: "b_cols"
832+
833+
a: Mapped[int] = mapped_column(primary_key=True)
834+
b: Mapped[str]
835+
836+
837+
class b_cols(B, TypedColumns):
838+
pass
839+
840+
841+
# table_b is a FromClause[b_cols], can call with just B since it
842+
# provides the __typed_cols__ annotation
843+
table_b = as_typed_table(B)
844+
845+
For proper typing integration :class:`_orm.MappedColumn` should be used
846+
to annotate the single columns, since it's a more specific annotation than
847+
the usual :class:`_orm.Mapped` used for ORM attributes.
848+
849+
"""
850+
return cls.__table__
851+
852+
756853
class DeclarativeBase(
757854
# Inspectable is used only by the mypy plugin
758855
inspection.Inspectable[InstanceState[Any]],

0 commit comments

Comments
 (0)