names.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342
  1. # ext/mypy/names.py
  2. # Copyright (C) 2021-2024 the SQLAlchemy authors and contributors
  3. # <see AUTHORS file>
  4. #
  5. # This module is part of SQLAlchemy and is released under
  6. # the MIT License: https://www.opensource.org/licenses/mit-license.php
  7. from __future__ import annotations
  8. from typing import Dict
  9. from typing import List
  10. from typing import Optional
  11. from typing import Set
  12. from typing import Tuple
  13. from typing import Union
  14. from mypy.nodes import ARG_POS
  15. from mypy.nodes import CallExpr
  16. from mypy.nodes import ClassDef
  17. from mypy.nodes import Decorator
  18. from mypy.nodes import Expression
  19. from mypy.nodes import FuncDef
  20. from mypy.nodes import MemberExpr
  21. from mypy.nodes import NameExpr
  22. from mypy.nodes import OverloadedFuncDef
  23. from mypy.nodes import SymbolNode
  24. from mypy.nodes import TypeAlias
  25. from mypy.nodes import TypeInfo
  26. from mypy.plugin import SemanticAnalyzerPluginInterface
  27. from mypy.types import CallableType
  28. from mypy.types import get_proper_type
  29. from mypy.types import Instance
  30. from mypy.types import UnboundType
  31. from ... import util
  32. COLUMN: int = util.symbol("COLUMN")
  33. RELATIONSHIP: int = util.symbol("RELATIONSHIP")
  34. REGISTRY: int = util.symbol("REGISTRY")
  35. COLUMN_PROPERTY: int = util.symbol("COLUMN_PROPERTY")
  36. TYPEENGINE: int = util.symbol("TYPEENGNE")
  37. MAPPED: int = util.symbol("MAPPED")
  38. DECLARATIVE_BASE: int = util.symbol("DECLARATIVE_BASE")
  39. DECLARATIVE_META: int = util.symbol("DECLARATIVE_META")
  40. MAPPED_DECORATOR: int = util.symbol("MAPPED_DECORATOR")
  41. SYNONYM_PROPERTY: int = util.symbol("SYNONYM_PROPERTY")
  42. COMPOSITE_PROPERTY: int = util.symbol("COMPOSITE_PROPERTY")
  43. DECLARED_ATTR: int = util.symbol("DECLARED_ATTR")
  44. MAPPER_PROPERTY: int = util.symbol("MAPPER_PROPERTY")
  45. AS_DECLARATIVE: int = util.symbol("AS_DECLARATIVE")
  46. AS_DECLARATIVE_BASE: int = util.symbol("AS_DECLARATIVE_BASE")
  47. DECLARATIVE_MIXIN: int = util.symbol("DECLARATIVE_MIXIN")
  48. QUERY_EXPRESSION: int = util.symbol("QUERY_EXPRESSION")
  49. # names that must succeed with mypy.api.named_type
  50. NAMED_TYPE_BUILTINS_OBJECT = "builtins.object"
  51. NAMED_TYPE_BUILTINS_STR = "builtins.str"
  52. NAMED_TYPE_BUILTINS_LIST = "builtins.list"
  53. NAMED_TYPE_SQLA_MAPPED = "sqlalchemy.orm.base.Mapped"
  54. _lookup: Dict[str, Tuple[int, Set[str]]] = {
  55. "Column": (
  56. COLUMN,
  57. {
  58. "sqlalchemy.sql.schema.Column",
  59. "sqlalchemy.sql.Column",
  60. },
  61. ),
  62. "Relationship": (
  63. RELATIONSHIP,
  64. {
  65. "sqlalchemy.orm.relationships.Relationship",
  66. "sqlalchemy.orm.relationships.RelationshipProperty",
  67. "sqlalchemy.orm.Relationship",
  68. "sqlalchemy.orm.RelationshipProperty",
  69. },
  70. ),
  71. "RelationshipProperty": (
  72. RELATIONSHIP,
  73. {
  74. "sqlalchemy.orm.relationships.Relationship",
  75. "sqlalchemy.orm.relationships.RelationshipProperty",
  76. "sqlalchemy.orm.Relationship",
  77. "sqlalchemy.orm.RelationshipProperty",
  78. },
  79. ),
  80. "registry": (
  81. REGISTRY,
  82. {
  83. "sqlalchemy.orm.decl_api.registry",
  84. "sqlalchemy.orm.registry",
  85. },
  86. ),
  87. "ColumnProperty": (
  88. COLUMN_PROPERTY,
  89. {
  90. "sqlalchemy.orm.properties.MappedSQLExpression",
  91. "sqlalchemy.orm.MappedSQLExpression",
  92. "sqlalchemy.orm.properties.ColumnProperty",
  93. "sqlalchemy.orm.ColumnProperty",
  94. },
  95. ),
  96. "MappedSQLExpression": (
  97. COLUMN_PROPERTY,
  98. {
  99. "sqlalchemy.orm.properties.MappedSQLExpression",
  100. "sqlalchemy.orm.MappedSQLExpression",
  101. "sqlalchemy.orm.properties.ColumnProperty",
  102. "sqlalchemy.orm.ColumnProperty",
  103. },
  104. ),
  105. "Synonym": (
  106. SYNONYM_PROPERTY,
  107. {
  108. "sqlalchemy.orm.descriptor_props.Synonym",
  109. "sqlalchemy.orm.Synonym",
  110. "sqlalchemy.orm.descriptor_props.SynonymProperty",
  111. "sqlalchemy.orm.SynonymProperty",
  112. },
  113. ),
  114. "SynonymProperty": (
  115. SYNONYM_PROPERTY,
  116. {
  117. "sqlalchemy.orm.descriptor_props.Synonym",
  118. "sqlalchemy.orm.Synonym",
  119. "sqlalchemy.orm.descriptor_props.SynonymProperty",
  120. "sqlalchemy.orm.SynonymProperty",
  121. },
  122. ),
  123. "Composite": (
  124. COMPOSITE_PROPERTY,
  125. {
  126. "sqlalchemy.orm.descriptor_props.Composite",
  127. "sqlalchemy.orm.Composite",
  128. "sqlalchemy.orm.descriptor_props.CompositeProperty",
  129. "sqlalchemy.orm.CompositeProperty",
  130. },
  131. ),
  132. "CompositeProperty": (
  133. COMPOSITE_PROPERTY,
  134. {
  135. "sqlalchemy.orm.descriptor_props.Composite",
  136. "sqlalchemy.orm.Composite",
  137. "sqlalchemy.orm.descriptor_props.CompositeProperty",
  138. "sqlalchemy.orm.CompositeProperty",
  139. },
  140. ),
  141. "MapperProperty": (
  142. MAPPER_PROPERTY,
  143. {
  144. "sqlalchemy.orm.interfaces.MapperProperty",
  145. "sqlalchemy.orm.MapperProperty",
  146. },
  147. ),
  148. "TypeEngine": (TYPEENGINE, {"sqlalchemy.sql.type_api.TypeEngine"}),
  149. "Mapped": (MAPPED, {NAMED_TYPE_SQLA_MAPPED}),
  150. "declarative_base": (
  151. DECLARATIVE_BASE,
  152. {
  153. "sqlalchemy.ext.declarative.declarative_base",
  154. "sqlalchemy.orm.declarative_base",
  155. "sqlalchemy.orm.decl_api.declarative_base",
  156. },
  157. ),
  158. "DeclarativeMeta": (
  159. DECLARATIVE_META,
  160. {
  161. "sqlalchemy.ext.declarative.DeclarativeMeta",
  162. "sqlalchemy.orm.DeclarativeMeta",
  163. "sqlalchemy.orm.decl_api.DeclarativeMeta",
  164. },
  165. ),
  166. "mapped": (
  167. MAPPED_DECORATOR,
  168. {
  169. "sqlalchemy.orm.decl_api.registry.mapped",
  170. "sqlalchemy.orm.registry.mapped",
  171. },
  172. ),
  173. "as_declarative": (
  174. AS_DECLARATIVE,
  175. {
  176. "sqlalchemy.ext.declarative.as_declarative",
  177. "sqlalchemy.orm.decl_api.as_declarative",
  178. "sqlalchemy.orm.as_declarative",
  179. },
  180. ),
  181. "as_declarative_base": (
  182. AS_DECLARATIVE_BASE,
  183. {
  184. "sqlalchemy.orm.decl_api.registry.as_declarative_base",
  185. "sqlalchemy.orm.registry.as_declarative_base",
  186. },
  187. ),
  188. "declared_attr": (
  189. DECLARED_ATTR,
  190. {
  191. "sqlalchemy.orm.decl_api.declared_attr",
  192. "sqlalchemy.orm.declared_attr",
  193. },
  194. ),
  195. "declarative_mixin": (
  196. DECLARATIVE_MIXIN,
  197. {
  198. "sqlalchemy.orm.decl_api.declarative_mixin",
  199. "sqlalchemy.orm.declarative_mixin",
  200. },
  201. ),
  202. "query_expression": (
  203. QUERY_EXPRESSION,
  204. {
  205. "sqlalchemy.orm.query_expression",
  206. "sqlalchemy.orm._orm_constructors.query_expression",
  207. },
  208. ),
  209. }
  210. def has_base_type_id(info: TypeInfo, type_id: int) -> bool:
  211. for mr in info.mro:
  212. check_type_id, fullnames = _lookup.get(mr.name, (None, None))
  213. if check_type_id == type_id:
  214. break
  215. else:
  216. return False
  217. if fullnames is None:
  218. return False
  219. return mr.fullname in fullnames
  220. def mro_has_id(mro: List[TypeInfo], type_id: int) -> bool:
  221. for mr in mro:
  222. check_type_id, fullnames = _lookup.get(mr.name, (None, None))
  223. if check_type_id == type_id:
  224. break
  225. else:
  226. return False
  227. if fullnames is None:
  228. return False
  229. return mr.fullname in fullnames
  230. def type_id_for_unbound_type(
  231. type_: UnboundType, cls: ClassDef, api: SemanticAnalyzerPluginInterface
  232. ) -> Optional[int]:
  233. sym = api.lookup_qualified(type_.name, type_)
  234. if sym is not None:
  235. if isinstance(sym.node, TypeAlias):
  236. target_type = get_proper_type(sym.node.target)
  237. if isinstance(target_type, Instance):
  238. return type_id_for_named_node(target_type.type)
  239. elif isinstance(sym.node, TypeInfo):
  240. return type_id_for_named_node(sym.node)
  241. return None
  242. def type_id_for_callee(callee: Expression) -> Optional[int]:
  243. if isinstance(callee, (MemberExpr, NameExpr)):
  244. if isinstance(callee.node, Decorator) and isinstance(
  245. callee.node.func, FuncDef
  246. ):
  247. if callee.node.func.type and isinstance(
  248. callee.node.func.type, CallableType
  249. ):
  250. ret_type = get_proper_type(callee.node.func.type.ret_type)
  251. if isinstance(ret_type, Instance):
  252. return type_id_for_fullname(ret_type.type.fullname)
  253. return None
  254. elif isinstance(callee.node, OverloadedFuncDef):
  255. if (
  256. callee.node.impl
  257. and callee.node.impl.type
  258. and isinstance(callee.node.impl.type, CallableType)
  259. ):
  260. ret_type = get_proper_type(callee.node.impl.type.ret_type)
  261. if isinstance(ret_type, Instance):
  262. return type_id_for_fullname(ret_type.type.fullname)
  263. return None
  264. elif isinstance(callee.node, FuncDef):
  265. if callee.node.type and isinstance(callee.node.type, CallableType):
  266. ret_type = get_proper_type(callee.node.type.ret_type)
  267. if isinstance(ret_type, Instance):
  268. return type_id_for_fullname(ret_type.type.fullname)
  269. return None
  270. elif isinstance(callee.node, TypeAlias):
  271. target_type = get_proper_type(callee.node.target)
  272. if isinstance(target_type, Instance):
  273. return type_id_for_fullname(target_type.type.fullname)
  274. elif isinstance(callee.node, TypeInfo):
  275. return type_id_for_named_node(callee)
  276. return None
  277. def type_id_for_named_node(
  278. node: Union[NameExpr, MemberExpr, SymbolNode]
  279. ) -> Optional[int]:
  280. type_id, fullnames = _lookup.get(node.name, (None, None))
  281. if type_id is None or fullnames is None:
  282. return None
  283. elif node.fullname in fullnames:
  284. return type_id
  285. else:
  286. return None
  287. def type_id_for_fullname(fullname: str) -> Optional[int]:
  288. tokens = fullname.split(".")
  289. immediate = tokens[-1]
  290. type_id, fullnames = _lookup.get(immediate, (None, None))
  291. if type_id is None or fullnames is None:
  292. return None
  293. elif fullname in fullnames:
  294. return type_id
  295. else:
  296. return None
  297. def expr_to_mapped_constructor(expr: Expression) -> CallExpr:
  298. column_descriptor = NameExpr("__sa_Mapped")
  299. column_descriptor.fullname = NAMED_TYPE_SQLA_MAPPED
  300. member_expr = MemberExpr(column_descriptor, "_empty_constructor")
  301. return CallExpr(
  302. member_expr,
  303. [expr],
  304. [ARG_POS],
  305. ["arg1"],
  306. )