names.py 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253
  1. # ext/mypy/names.py
  2. # Copyright (C) 2021 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 typing import Dict
  8. from typing import List
  9. from typing import Optional
  10. from typing import Set
  11. from typing import Tuple
  12. from typing import Union
  13. from mypy.nodes import ClassDef
  14. from mypy.nodes import Expression
  15. from mypy.nodes import FuncDef
  16. from mypy.nodes import MemberExpr
  17. from mypy.nodes import NameExpr
  18. from mypy.nodes import SymbolNode
  19. from mypy.nodes import TypeAlias
  20. from mypy.nodes import TypeInfo
  21. from mypy.plugin import SemanticAnalyzerPluginInterface
  22. from mypy.types import CallableType
  23. from mypy.types import get_proper_type
  24. from mypy.types import Instance
  25. from mypy.types import UnboundType
  26. from ... import util
  27. COLUMN: int = util.symbol("COLUMN") # type: ignore
  28. RELATIONSHIP: int = util.symbol("RELATIONSHIP") # type: ignore
  29. REGISTRY: int = util.symbol("REGISTRY") # type: ignore
  30. COLUMN_PROPERTY: int = util.symbol("COLUMN_PROPERTY") # type: ignore
  31. TYPEENGINE: int = util.symbol("TYPEENGNE") # type: ignore
  32. MAPPED: int = util.symbol("MAPPED") # type: ignore
  33. DECLARATIVE_BASE: int = util.symbol("DECLARATIVE_BASE") # type: ignore
  34. DECLARATIVE_META: int = util.symbol("DECLARATIVE_META") # type: ignore
  35. MAPPED_DECORATOR: int = util.symbol("MAPPED_DECORATOR") # type: ignore
  36. COLUMN_PROPERTY: int = util.symbol("COLUMN_PROPERTY") # type: ignore
  37. SYNONYM_PROPERTY: int = util.symbol("SYNONYM_PROPERTY") # type: ignore
  38. COMPOSITE_PROPERTY: int = util.symbol("COMPOSITE_PROPERTY") # type: ignore
  39. DECLARED_ATTR: int = util.symbol("DECLARED_ATTR") # type: ignore
  40. MAPPER_PROPERTY: int = util.symbol("MAPPER_PROPERTY") # type: ignore
  41. AS_DECLARATIVE: int = util.symbol("AS_DECLARATIVE") # type: ignore
  42. AS_DECLARATIVE_BASE: int = util.symbol("AS_DECLARATIVE_BASE") # type: ignore
  43. DECLARATIVE_MIXIN: int = util.symbol("DECLARATIVE_MIXIN") # type: ignore
  44. QUERY_EXPRESSION: int = util.symbol("QUERY_EXPRESSION") # type: ignore
  45. # names that must succeed with mypy.api.named_type
  46. NAMED_TYPE_BUILTINS_OBJECT = "builtins.object"
  47. NAMED_TYPE_BUILTINS_STR = "builtins.str"
  48. NAMED_TYPE_BUILTINS_LIST = "builtins.list"
  49. NAMED_TYPE_SQLA_MAPPED = "sqlalchemy.orm.attributes.Mapped"
  50. _lookup: Dict[str, Tuple[int, Set[str]]] = {
  51. "Column": (
  52. COLUMN,
  53. {
  54. "sqlalchemy.sql.schema.Column",
  55. "sqlalchemy.sql.Column",
  56. },
  57. ),
  58. "RelationshipProperty": (
  59. RELATIONSHIP,
  60. {
  61. "sqlalchemy.orm.relationships.RelationshipProperty",
  62. "sqlalchemy.orm.RelationshipProperty",
  63. },
  64. ),
  65. "registry": (
  66. REGISTRY,
  67. {
  68. "sqlalchemy.orm.decl_api.registry",
  69. "sqlalchemy.orm.registry",
  70. },
  71. ),
  72. "ColumnProperty": (
  73. COLUMN_PROPERTY,
  74. {
  75. "sqlalchemy.orm.properties.ColumnProperty",
  76. "sqlalchemy.orm.ColumnProperty",
  77. },
  78. ),
  79. "SynonymProperty": (
  80. SYNONYM_PROPERTY,
  81. {
  82. "sqlalchemy.orm.descriptor_props.SynonymProperty",
  83. "sqlalchemy.orm.SynonymProperty",
  84. },
  85. ),
  86. "CompositeProperty": (
  87. COMPOSITE_PROPERTY,
  88. {
  89. "sqlalchemy.orm.descriptor_props.CompositeProperty",
  90. "sqlalchemy.orm.CompositeProperty",
  91. },
  92. ),
  93. "MapperProperty": (
  94. MAPPER_PROPERTY,
  95. {
  96. "sqlalchemy.orm.interfaces.MapperProperty",
  97. "sqlalchemy.orm.MapperProperty",
  98. },
  99. ),
  100. "TypeEngine": (TYPEENGINE, {"sqlalchemy.sql.type_api.TypeEngine"}),
  101. "Mapped": (MAPPED, {"sqlalchemy.orm.attributes.Mapped"}),
  102. "declarative_base": (
  103. DECLARATIVE_BASE,
  104. {
  105. "sqlalchemy.ext.declarative.declarative_base",
  106. "sqlalchemy.orm.declarative_base",
  107. "sqlalchemy.orm.decl_api.declarative_base",
  108. },
  109. ),
  110. "DeclarativeMeta": (
  111. DECLARATIVE_META,
  112. {
  113. "sqlalchemy.ext.declarative.DeclarativeMeta",
  114. "sqlalchemy.orm.DeclarativeMeta",
  115. "sqlalchemy.orm.decl_api.DeclarativeMeta",
  116. },
  117. ),
  118. "mapped": (
  119. MAPPED_DECORATOR,
  120. {
  121. "sqlalchemy.orm.decl_api.registry.mapped",
  122. "sqlalchemy.orm.registry.mapped",
  123. },
  124. ),
  125. "as_declarative": (
  126. AS_DECLARATIVE,
  127. {
  128. "sqlalchemy.ext.declarative.as_declarative",
  129. "sqlalchemy.orm.decl_api.as_declarative",
  130. "sqlalchemy.orm.as_declarative",
  131. },
  132. ),
  133. "as_declarative_base": (
  134. AS_DECLARATIVE_BASE,
  135. {
  136. "sqlalchemy.orm.decl_api.registry.as_declarative_base",
  137. "sqlalchemy.orm.registry.as_declarative_base",
  138. },
  139. ),
  140. "declared_attr": (
  141. DECLARED_ATTR,
  142. {
  143. "sqlalchemy.orm.decl_api.declared_attr",
  144. "sqlalchemy.orm.declared_attr",
  145. },
  146. ),
  147. "declarative_mixin": (
  148. DECLARATIVE_MIXIN,
  149. {
  150. "sqlalchemy.orm.decl_api.declarative_mixin",
  151. "sqlalchemy.orm.declarative_mixin",
  152. },
  153. ),
  154. "query_expression": (
  155. QUERY_EXPRESSION,
  156. {"sqlalchemy.orm.query_expression"},
  157. ),
  158. }
  159. def has_base_type_id(info: TypeInfo, type_id: int) -> bool:
  160. for mr in info.mro:
  161. check_type_id, fullnames = _lookup.get(mr.name, (None, None))
  162. if check_type_id == type_id:
  163. break
  164. else:
  165. return False
  166. if fullnames is None:
  167. return False
  168. return mr.fullname in fullnames
  169. def mro_has_id(mro: List[TypeInfo], type_id: int) -> bool:
  170. for mr in mro:
  171. check_type_id, fullnames = _lookup.get(mr.name, (None, None))
  172. if check_type_id == type_id:
  173. break
  174. else:
  175. return False
  176. if fullnames is None:
  177. return False
  178. return mr.fullname in fullnames
  179. def type_id_for_unbound_type(
  180. type_: UnboundType, cls: ClassDef, api: SemanticAnalyzerPluginInterface
  181. ) -> Optional[int]:
  182. sym = api.lookup_qualified(type_.name, type_)
  183. if sym is not None:
  184. if isinstance(sym.node, TypeAlias):
  185. target_type = get_proper_type(sym.node.target)
  186. if isinstance(target_type, Instance):
  187. return type_id_for_named_node(target_type.type)
  188. elif isinstance(sym.node, TypeInfo):
  189. return type_id_for_named_node(sym.node)
  190. return None
  191. def type_id_for_callee(callee: Expression) -> Optional[int]:
  192. if isinstance(callee, (MemberExpr, NameExpr)):
  193. if isinstance(callee.node, FuncDef):
  194. if callee.node.type and isinstance(callee.node.type, CallableType):
  195. ret_type = get_proper_type(callee.node.type.ret_type)
  196. if isinstance(ret_type, Instance):
  197. return type_id_for_fullname(ret_type.type.fullname)
  198. return None
  199. elif isinstance(callee.node, TypeAlias):
  200. target_type = get_proper_type(callee.node.target)
  201. if isinstance(target_type, Instance):
  202. return type_id_for_fullname(target_type.type.fullname)
  203. elif isinstance(callee.node, TypeInfo):
  204. return type_id_for_named_node(callee)
  205. return None
  206. def type_id_for_named_node(
  207. node: Union[NameExpr, MemberExpr, SymbolNode]
  208. ) -> Optional[int]:
  209. type_id, fullnames = _lookup.get(node.name, (None, None))
  210. if type_id is None or fullnames is None:
  211. return None
  212. elif node.fullname in fullnames:
  213. return type_id
  214. else:
  215. return None
  216. def type_id_for_fullname(fullname: str) -> Optional[int]:
  217. tokens = fullname.split(".")
  218. immediate = tokens[-1]
  219. type_id, fullnames = _lookup.get(immediate, (None, None))
  220. if type_id is None or fullnames is None:
  221. return None
  222. elif fullname in fullnames:
  223. return type_id
  224. else:
  225. return None