def _apply_type_to_mapped_statement(
api: SemanticAnalyzerPluginInterface,
stmt: AssignmentStmt,
lvalue: NameExpr,
left_hand_explicit_type: Optional[Union[Instance, UnionType]],
python_type_for_type: Union[Instance, UnionType],
) -> None:
"""Apply the Mapped[<type>] annotation and right hand object to a
declarative assignment statement.
This converts a Python declarative class statement such as::
class User(Base):
# ...
attrname = Column(Integer)
To one that describes the final Python behavior to Mypy::
class User(Base):
# ...
attrname : Mapped[Optional[int]] = <meaningless temp node>
"""
descriptor = api.modules["sqlalchemy.orm.attributes"].names["Mapped"]
left_node = lvalue.node
inst = Instance(descriptor.node, [python_type_for_type])
if left_hand_explicit_type is not None:
left_node.type = Instance(descriptor.node, [left_hand_explicit_type])
else:
lvalue.is_inferred_def = False
left_node.type = inst
# so to have it skip the right side totally, we can do this:
# stmt.rvalue = TempNode(AnyType(TypeOfAny.special_form))
# however, if we instead manufacture a new node that uses the old
# one, then we can still get type checking for the call itself,
# e.g. the Column, relationship() call, etc.
# rewrite the node as:
# <attr> : Mapped[<typ>] =
# _sa_Mapped._empty_constructor(<original CallExpr from rvalue>)
# the original right-hand side is maintained so it gets type checked
# internally
api.add_symbol_table_node("_sa_Mapped", descriptor)
column_descriptor = nodes.NameExpr("_sa_Mapped")
column_descriptor.fullname = "sqlalchemy.orm.Mapped"
mm = nodes.MemberExpr(column_descriptor, "_empty_constructor")
orig_call_expr = stmt.rvalue
stmt.rvalue = CallExpr(
mm,
[orig_call_expr],
[nodes.ARG_POS],
["arg1"],
)
def add_model_astuple(ctx) -> None:
"""Add model astuple method."""
bool_type = ctx.api.builtin_type('builtins.bool')
tuple_type = ctx.api.builtin_type('builtins.tuple')
var = nodes.Var('recurse', bool_type)
recurse = nodes.Argument(variable=var,
type_annotation=bool_type,
initializer=nodes.NameExpr('True'),
kind=nodes.ARG_NAMED_OPT)
common.add_method(ctx, 'astuple', [recurse], tuple_type)
def _scan_declarative_decorator_stmt(
cls: ClassDef,
api: SemanticAnalyzerPluginInterface,
stmt: Decorator,
cls_metadata: util.DeclClassApplied,
) -> None:
"""Extract mapping information from a @declared_attr in a declarative
class.
E.g.::
@reg.mapped
class MyClass:
# ...
@declared_attr
def updated_at(cls) -> Column[DateTime]:
return Column(DateTime)
Will resolve in mypy as::
@reg.mapped
class MyClass:
# ...
updated_at: Mapped[Optional[datetime.datetime]]
"""
for dec in stmt.decorators:
if (
isinstance(dec, (NameExpr, MemberExpr, SymbolNode))
and names._type_id_for_named_node(dec) is names.DECLARED_ATTR
):
break
else:
return
dec_index = cls.defs.body.index(stmt)
left_hand_explicit_type: Optional[ProperType] = None
if isinstance(stmt.func.type, CallableType):
func_type = stmt.func.type.ret_type
if isinstance(func_type, UnboundType):
type_id = names._type_id_for_unbound_type(func_type, cls, api)
else:
# this does not seem to occur unless the type argument is
# incorrect
return
if (
type_id
in {
names.MAPPED,
names.RELATIONSHIP,
names.COMPOSITE_PROPERTY,
names.MAPPER_PROPERTY,
names.SYNONYM_PROPERTY,
names.COLUMN_PROPERTY,
}
and func_type.args
):
left_hand_explicit_type = get_proper_type(func_type.args[0])
elif type_id is names.COLUMN and func_type.args:
typeengine_arg = func_type.args[0]
if isinstance(typeengine_arg, UnboundType):
sym = api.lookup_qualified(typeengine_arg.name, typeengine_arg)
if sym is not None and isinstance(sym.node, TypeInfo):
if names._has_base_type_id(sym.node, names.TYPEENGINE):
left_hand_explicit_type = UnionType(
[
infer._extract_python_type_from_typeengine(
api, sym.node, []
),
NoneType(),
]
)
else:
util.fail(
api,
"Column type should be a TypeEngine "
"subclass not '{}'".format(sym.node.fullname),
func_type,
)
if left_hand_explicit_type is None:
# no type on the decorated function. our option here is to
# dig into the function body and get the return type, but they
# should just have an annotation.
msg = (
"Can't infer type from @declared_attr on function '{}'; "
"please specify a return type from this function that is "
"one of: Mapped[<python type>], relationship[<target class>], "
"Column[<TypeEngine>], MapperProperty[<python type>]"
)
util.fail(api, msg.format(stmt.var.name), stmt)
left_hand_explicit_type = AnyType(TypeOfAny.special_form)
left_node = NameExpr(stmt.var.name)
left_node.node = stmt.var
# totally feeling around in the dark here as I don't totally understand
# the significance of UnboundType. It seems to be something that is
# not going to do what's expected when it is applied as the type of
# an AssignmentStatement. So do a feeling-around-in-the-dark version
# of converting it to the regular Instance/TypeInfo/UnionType structures
# we see everywhere else.
if isinstance(left_hand_explicit_type, UnboundType):
left_hand_explicit_type = get_proper_type(
util._unbound_to_instance(api, left_hand_explicit_type)
)
left_node.node.type = api.named_type(
"__sa_Mapped", [left_hand_explicit_type]
)
# this will ignore the rvalue entirely
# rvalue = TempNode(AnyType(TypeOfAny.special_form))
# rewrite the node as:
# <attr> : Mapped[<typ>] =
# _sa_Mapped._empty_constructor(lambda: <function body>)
# the function body is maintained so it gets type checked internally
column_descriptor = nodes.NameExpr("__sa_Mapped")
column_descriptor.fullname = "sqlalchemy.orm.attributes.Mapped"
mm = nodes.MemberExpr(column_descriptor, "_empty_constructor")
arg = nodes.LambdaExpr(stmt.func.arguments, stmt.func.body)
rvalue = CallExpr(
mm,
[arg],
[nodes.ARG_POS],
["arg1"],
)
new_stmt = AssignmentStmt([left_node], rvalue)
new_stmt.type = left_node.node.type
cls_metadata.mapped_attr_names.append(
(left_node.name, left_hand_explicit_type)
)
cls.defs.body[dec_index] = new_stmt
def transform_enum_type(context: mypy.plugin.AnalyzeTypeContext) -> types.Type:
'''
This is registered as a handler of the :code:`get_type_analyze_hook` hook
See https://mypy.readthedocs.io/en/latest/extending_mypy.html#current-list-of-plugin-hooks for
more information about hooks.
This will be the first hook called in this plugin.
It allows us to change or alter type definitions as mypy sees them.
This is needed because our Enum uses :code:`EnumMeta` class which defines :code:`__new__`
and it shadows everything for mypy. So without this callback our type
visible by mypy would look like following in :code:`transform_enum_class_def`
hook handler. (Please read comment in :code:`transform_enum_class_def`).
::
ClassDef:3(
Color
FallbackToAny
AssignmentStmt:4(
NameExpr(RED [m])
IntExpr(1)
builtins.int)
AssignmentStmt:5(
NameExpr(GREEN [m])
IntExpr(2)
builtins.int)
AssignmentStmt:6(
NameExpr(BLUE [m])
IntExpr(3)
builtins.int))
TypeInfo(
Name(main.Color)
Bases(builtins.object)
Mro(main.Color, builtins.object)
Names(
BLUE (builtins.int)
GREEN (builtins.int)
RED (builtins.int)))
Please check docs to :code:`transform_enum_class_def` method
where you can se what these types look like when this hook
is used.
For example you can see that :code:`TypeInfo` is completely missing
- inheritance to Enum type
- metaclass definition
So it would be possible to change attribute types without this hook
but we have to add a metaclass for :code:`__iter__` and :code:`__getitem__` definitions.
To do this we have to describe what :code:`Enum` and :code:`EnumMeta` look like.
We did not find how to tell mypy to give us these definitions
so we described them manually in this part, and then properly
added inheritance to any :code:`Enum` inherited class to our fake :code:`Enum`.
'''
# Find references to some builtins which are often used
type_type = context.api.named_type('builtins.type')
object_type = context.api.named_type('builtins.object')
str_type = context.api.named_type('builtins.str')
bool_type = context.api.named_type('builtins.bool')
# Define meta class in fake module :code:`_fastenum`
# This is roughly equivalent to:
#
# class EnumMeta(builtins.type): pass
#
# in module :code:`_fastenum`.
# We have to define two things: :code:`ClassDef` an AST node and it's type definition using :code:`TypeInfo`
# :code:`ClassDef` defines only class syntactically all attributes and such will be defined in :code:`TypeInfo`
meta_cls = nodes.ClassDef('EnumMeta', nodes.Block([nodes.PassStmt()]), [],
[type_type])
meta_cls.fullname = '_fastenum.EnumMeta'
meta_info = nodes.TypeInfo(nodes.SymbolTable(), meta_cls, '_fastenum')
# We have to define inheritance again, mypy :code:`ClassDef` and :code:`TypeInfo`
# won't automatically share this information so we have to tell it again it is inherited from :code:`builtins.type`
meta_info.bases = [type_type]
# Last thing to get everything working is to define mro (method resolution order)
# without correctly specifying this mypy won't complain but it wont see any method or attributes
# defined in parents or even class itself.
# So we have to define class itself as :code:`meta_info` and **all of it's parents** (even indirect one)
# (this is not working in transitional fashion)
meta_info.mro = [meta_info, type_type.type, object_type.type]
# Define Enum class which is using EnumMeta as its metaclass in the fake :code:`_fastenum` module
# This is very similar to the previous definition and it is roughly equivalent to:
#
# class Enum(metaclass = EnumMeta): pass
#
# Notice that we still define :code:`builtins.object` as it's parent
# even if we don't have to do it in Python3, but we have to do it here!
enum_cls = nodes.ClassDef('Enum', nodes.Block([nodes.PassStmt()]), [],
[object_type], nodes.NameExpr('EnumMeta'))
enum_cls.fullname = '_fastenum.Enum'
enum_info = nodes.TypeInfo(nodes.SymbolTable(), enum_cls, '_fastenum')
# Same as before we have to define all parents (even :code:`builtins.object`)
enum_info.bases = [object_type]
enum_info.mro = [enum_info, object_type.type]
# New things in here are that we have to define the metaclass again in info
# I don't know why we have to define it on :code:`metaclass_type` and :code:`declared_metaclass`
# at the same time but mypy requires it that way, otherwise it ignores that metaclass
# and does not complain at all.
enum_info.metaclass_type = types.Instance(meta_info, [])
enum_info.declared_metaclass = types.Instance(meta_info, [])
# Add the attribute ``value`` to enum instances.
# Don't be scared by :code:`TypeOfAny`. mypy just has multiple types of Any.
# See :code:`TypeOfAny` definition (it is an enum with comments).
value_attribute = nodes.Var('value',
types.AnyType(types.TypeOfAny.explicit))
value_attribute.is_initialized_in_class = False
# As before we have to link our variable back to our class.
value_attribute.info = enum_info
enum_info.names['value'] = nodes.SymbolTableNode(nodes.MDEF,
value_attribute,
plugin_generated=True)
# Add the attribute ``name``` to enum instances.
value_attribute = nodes.Var('name', str_type)
value_attribute.is_initialized_in_class = False
value_attribute.info = enum_info
enum_info.names['name'] = nodes.SymbolTableNode(nodes.MDEF,
value_attribute,
plugin_generated=True)
#
# So after these few lines we end up with something like:
#
# module `_fastenum`:
#
# class EnumMeta(builtins.type): pass
#
# class Enum(metaclass = EnumMeta):
# name: str
# value: Any
#
# Prepare TypeVar, all these lines are just:
# _EnumMetaType = TypeVar('_EnumMetaType', bound = 'EnumMeta')
# We just have to describe expressions and definitions separately for mypy
meta_enum_instance = types.Instance(meta_info, [])
self_tvar_expr = nodes.TypeVarExpr(
'_EnumMetaType', f'{meta_info.fullname()}._EnumMetaType', [],
meta_enum_instance)
meta_info.names['_EnumMetaType'] = nodes.SymbolTableNode(
nodes.MDEF, self_tvar_expr)
self_tvar_def = types.TypeVarDef('_EnumMetaType',
f'{meta_info.fullname()}._EnumMetaType',
-1, [], meta_enum_instance)
self_tvar_type = types.TypeVarType(self_tvar_def)
# Define base __iter__ and __next__ for our meta class and use TypeVar `_EnumMetaType` as its return value
# so we can say its return value is bound to all children.
# See more comments about the definition in:
# - `transform_enum_class_def` handler
# - and `_define_method` docs + comments
_define_method(context, meta_info, context.type.name, '__iter__', [],
self_tvar_type)
_define_method(context, meta_info, context.type.name, '__next__', [],
self_tvar_type)
# Same way with __getitem__
_define_method(context, meta_info, context.type.name, '__getitem__', [
nodes.Argument(nodes.Var('cls', meta_enum_instance),
meta_enum_instance, None, nodes.ARG_POS),
nodes.Argument(nodes.Var('key', str_type), str_type, None,
nodes.ARG_POS),
], self_tvar_type)
# We also have to support constructor interface of enum, so when someone calls Enum('value').
# This is simply done by adding the `__init__` method with two arguments (self, value).
enum_instance = types.Instance(enum_info, [])
any_type = types.AnyType(types.TypeOfAny.explicit)
_define_method(context, enum_info, context.type.name, '__init__', [
nodes.Argument(nodes.Var('self', enum_instance), enum_instance, None,
nodes.ARG_POS),
nodes.Argument(nodes.Var('value', any_type), any_type, None,
nodes.ARG_POS),
], types.NoneTyp())
# Because enums can be used even in comparison expression like `A > B`
# we have to support these methods in our fake enum class too.
def def_bool_method(name: str) -> None:
_define_method(context, enum_info, context.type.name, name, [
nodes.Argument(nodes.Var('self', enum_instance), enum_instance,
None, nodes.ARG_POS),
nodes.Argument(nodes.Var('other', enum_instance), enum_instance,
None, nodes.ARG_POS),
], bool_type)
for name in ('le', 'eq', 'ne', 'ge', 'gt'):
def_bool_method(f'__{name}__')
#
# After all this we end up with:
#
# module `_fastenum`:
#
# class EnumMeta(builtins.type):
# _EnumMetaType = TypeVar('_EnumMetaType', bound = 'EnumMeta')
#
# def __iter__() -> _EnumMetaType: pass
# def __next__() -> _EnumMetaType: pass
# def __getitem__(cls: 'EnumMeta', key: str) -> 'EnumMeta': pass
#
#
# class Enum(metaclass = EnumMeta):
# name: str
# value: Any
#
# def __init__(self, value: Any) -> None: pass
#
# def __le__(self, other: Enum) -> bool: pass
# def __eq__(self, other: Enum) -> bool: pass
# def __ne__(self, other: Enum) -> bool: pass
# def __ge__(self, other: Enum) -> bool: pass
# def __gt__(self, other: Enum) -> bool: pass
#
# And we have to return new type for our `Enum` class which will be our new `Enum`
return types.Instance(enum_info, [])
