def _scan_declarative_decorator_stmt(
cls: ClassDef,
api: SemanticAnalyzerPluginInterface,
stmt: Decorator,
cls_metadata: DeclClassApplied,
):
"""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 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 = None
if stmt.func.type is not None:
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 = 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(typeengine_arg.name, typeengine_arg)
if sym is not None and names._mro_has_id(
sym.node.mro, names.TYPEENGINE):
left_hand_explicit_type = UnionType([
_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)
descriptor = api.lookup("__sa_Mapped", cls)
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 = util._unbound_to_instance(
api, left_hand_explicit_type)
left_node.node.type = Instance(descriptor.node, [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.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 visit_Name(self, n: ast27.Name) -> NameExpr:
return NameExpr(n.id)
def visit_Name(self, n: Name) -> NameExpr:
e = NameExpr(n.id)
return self.set_line(e, n)
def visit_Name(self, n: ast35.Name) -> Node:
return NameExpr(n.id)
def duplicate_name(self, node: NameExpr) -> NameExpr:
# This method is used when the transform result must be a NameExpr.
# visit_name_expr() is used when there is no such restriction.
new = NameExpr(node.name)
self.copy_ref(new, node)
return new
def parse_key_typeddict_fields(attrs_expr: DictExpr) -> Tuple[List[str], List[Type], Set[str]]:
fields = []
types = []
required_fields = set()
for field_name_expr, field_type_expr in attrs_expr.items:
if isinstance(field_name_expr, StrExpr):
fields.append(field_name_expr.value)
required_fields.add(field_name_expr.value)
elif isinstance(field_name_expr, CallExpr):
fields.append(field_name_expr.args[0].value)
required_expr = field_name_expr.args[1] if len(field_name_expr.args) == 2 else NameExpr('builtins.True')
if required_expr.fullname == 'builtins.True':
required_fields.add(field_name_expr.args[0].value)
else:
raise UnsupportedKeyTypeError(str(type(field_name_expr)))
types.append(expr_to_unanalyzed_type(field_type_expr))
return fields, types, required_fields
def visit_NameConstant(self, n: ast35.NameConstant) -> Node:
return NameExpr(str(n.value))
def add_class_method(
ctx: ClassDefContext,
name: str,
args: List[Argument],
return_type: Type,
self_type: Optional[Type] = None,
tvar_def: Optional[TypeVarDef] = None,
) -> None:
"""Adds a new class method to a class.
"""
info = ctx.cls.info
# First remove any previously generated methods with the same name
# to avoid clashes and problems in the semantic analyzer.
if name in info.names:
sym = info.names[name]
if sym.plugin_generated and isinstance(sym.node, FuncDef):
ctx.cls.defs.body.remove(sym.node)
self_type = self_type or fill_typevars(info)
function_type = ctx.api.named_type('__builtins__.function')
args = [
Argument(Var('_cls'), TypeType.make_normalized(self_type), None,
ARG_POS)
] + args
arg_types, arg_names, arg_kinds = [], [], []
for arg in args:
assert arg.type_annotation, 'All arguments must be fully typed.'
arg_types.append(arg.type_annotation)
arg_names.append(arg.variable.name)
arg_kinds.append(arg.kind)
signature = CallableType(arg_types, arg_kinds, arg_names, return_type,
function_type)
if tvar_def:
signature.variables = [tvar_def]
func = FuncDef(name, args, Block([PassStmt()]))
func.info = info
func.type = set_callable_name(signature, func)
func.is_class = True
func._fullname = info.fullname + '.' + name
func.line = info.line
# NOTE: we would like the plugin generated node to dominate, but we still
# need to keep any existing definitions so they get semantically analyzed.
if name in info.names:
# Get a nice unique name instead.
r_name = get_unique_redefinition_name(name, info.names)
info.names[r_name] = info.names[name]
func.is_decorated = True
v = Var(name, func.type)
v.info = info
v._fullname = func._fullname
v.is_classmethod = True
dec = Decorator(func, [NameExpr('classmethod')], v)
dec.line = info.line
sym = SymbolTableNode(MDEF, dec)
sym.plugin_generated = True
info.names[name] = sym
info.defn.defs.body.append(func)
def visit_Name(self, n):
return NameExpr(n.id)
def visit_NameConstant(self, n):
return NameExpr(str(n.value))
def visit_NameConstant(self, n: NameConstant) -> NameExpr:
return NameExpr(str(n.value))
def add_method(
ctx: ClassDefContext,
name: str,
args: List[Argument],
return_type: Type,
self_type: Optional[Type] = None,
tvar_def: Optional[TypeVarDef] = None,
is_classmethod: bool = False,
is_new: bool = False,
# is_staticmethod: bool = False,
) -> None:
"""
Adds a new method to a class.
This can be dropped if/when https://github.com/python/mypy/issues/7301 is merged
"""
info = ctx.cls.info
# First remove any previously generated methods with the same name
# to avoid clashes and problems in the semantic analyzer.
if name in info.names:
sym = info.names[name]
if sym.plugin_generated and isinstance(sym.node, FuncDef):
ctx.cls.defs.body.remove(sym.node)
self_type = self_type or fill_typevars(info)
if is_classmethod or is_new:
first = [
Argument(Var('_cls'), TypeType.make_normalized(self_type), None,
ARG_POS)
]
# elif is_staticmethod:
# first = []
else:
self_type = self_type or fill_typevars(info)
first = [Argument(Var('__pydantic_self__'), self_type, None, ARG_POS)]
args = first + args
arg_types, arg_names, arg_kinds = [], [], []
for arg in args:
assert arg.type_annotation, 'All arguments must be fully typed.'
arg_types.append(arg.type_annotation)
arg_names.append(get_name(arg.variable))
arg_kinds.append(arg.kind)
function_type = ctx.api.named_type(f'{BUILTINS_NAME}.function')
signature = CallableType(arg_types, arg_kinds, arg_names, return_type,
function_type)
if tvar_def:
signature.variables = [tvar_def]
func = FuncDef(name, args, Block([PassStmt()]))
func.info = info
func.type = set_callable_name(signature, func)
func.is_class = is_classmethod
# func.is_static = is_staticmethod
func._fullname = get_fullname(info) + '.' + name
func.line = info.line
# NOTE: we would like the plugin generated node to dominate, but we still
# need to keep any existing definitions so they get semantically analyzed.
if name in info.names:
# Get a nice unique name instead.
r_name = get_unique_redefinition_name(name, info.names)
info.names[r_name] = info.names[name]
if is_classmethod: # or is_staticmethod:
func.is_decorated = True
v = Var(name, func.type)
v.info = info
v._fullname = func._fullname
# if is_classmethod:
v.is_classmethod = True
dec = Decorator(func, [NameExpr('classmethod')], v)
# else:
# v.is_staticmethod = True
# dec = Decorator(func, [NameExpr('staticmethod')], v)
dec.line = info.line
sym = SymbolTableNode(MDEF, dec)
else:
sym = SymbolTableNode(MDEF, func)
sym.plugin_generated = True
info.names[name] = sym
info.defn.defs.body.append(func)
