From 9470f5ccb5a934ade4064152b636f7f276755c38 Mon Sep 17 00:00:00 2001 From: Mehdi Date: Sat, 9 Feb 2019 17:12:10 +0100 Subject: [PATCH 1/7] wip --- sqlmypy.py | 37 ++++++++++++++++++++++++++++++++----- 1 file changed, 32 insertions(+), 5 deletions(-) diff --git a/sqlmypy.py b/sqlmypy.py index 9b66b7f..621efe1 100644 --- a/sqlmypy.py +++ b/sqlmypy.py @@ -3,10 +3,10 @@ Plugin, FunctionContext, ClassDefContext, DynamicClassDefContext, SemanticAnalyzerPluginInterface ) -from mypy.plugins.common import add_method +from mypy.plugins.common import add_method, _get_argument from mypy.nodes import ( NameExpr, Expression, StrExpr, TypeInfo, ClassDef, Block, SymbolTable, SymbolTableNode, GDEF, - Argument, Var, ARG_STAR2, MDEF, TupleExpr, RefExpr + Argument, Var, ARG_STAR2, MDEF, TupleExpr, RefExpr, AssignmentStmt, CallExpr, MemberExpr ) from mypy.types import ( UnionType, NoneTyp, Instance, Type, AnyType, TypeOfAny, UninhabitedType, CallableType @@ -25,6 +25,7 @@ COLUMN_ELEMENT_NAME = 'sqlalchemy.sql.elements.ColumnElement' # type: Final GROUPING_NAME = 'sqlalchemy.sql.elements.Grouping' # type: Final RELATIONSHIP_NAME = 'sqlalchemy.orm.relationships.RelationshipProperty' # type: Final +FOREIGN_KEY_NAME = 'sqlalchemy.sql.schema.ForeignKey' # type: Final def is_declarative(info: TypeInfo) -> bool: @@ -65,17 +66,17 @@ def get_function_hook(self, fullname: str) -> Optional[Callable[[FunctionContext return model_hook return None - def get_dynamic_class_hook(self, fullname: str) -> CB[DynamicClassDefContext]: + def get_dynamic_class_hook(self, fullname: str): if fullname == 'sqlalchemy.ext.declarative.api.declarative_base': return decl_info_hook return None - def get_class_decorator_hook(self, fullname: str) -> CB[ClassDefContext]: + def get_class_decorator_hook(self, fullname: str): if fullname == 'sqlalchemy.ext.declarative.api.as_declarative': return decl_deco_hook return None - def get_base_class_hook(self, fullname: str) -> CB[ClassDefContext]: + def get_base_class_hook(self, fullname: str): sym = self.lookup_fully_qualified(fullname) if sym and isinstance(sym.node, TypeInfo): if is_declarative(sym.node): @@ -110,6 +111,32 @@ def add_model_init_hook(ctx: ClassDefContext) -> None: add_method(ctx, '__init__', [kw_arg], NoneTyp()) ctx.cls.info.metadata.setdefault('sqlalchemy', {})['generated_init'] = True + for stmt in ctx.cls.defs.body: + if not (isinstance(stmt, AssignmentStmt) and len(stmt.lvalues) == 1 and isinstance(stmt.lvalues[0], NameExpr)): + continue + + if stmt.lvalues[0].name == "__tablename__" and isinstance(stmt.rvalue, StrExpr): + ctx.cls.info.metadata.setdefault('sqlalchemy', {})['tablename'] = stmt.rvalue.value + + if isinstance(stmt.rvalue, CallExpr) and stmt.rvalue.callee.fullname == COLUMN_NAME: + colname = stmt.lvalues[0].name + has_explicit_colname = stmt.rvalue + ctx.cls.info.metadata.setdefault('sqlalchemy', {}).setdefault('columns', []).append(colname) + for arg in stmt.rvalue.args: + if isinstance(arg, CallExpr) and arg.callee.fullname == FOREIGN_KEY_NAME and len(arg.args) >= 1: + fk = arg.args[0] + if isinstance(fk, StrExpr): + *_, parent_table, parent_col = fk.value.split(".") + ctx.cls.info.metadata.setdefault('sqlalchemy', {}).setdefault('foreign_keys', {})[colname] = { + "column": parent_col, + "table": parent_table + } + elif isinstance(fk, MemberExpr): + ctx.cls.info.metadata.setdefault('sqlalchemy', {}).setdefault('foreign_keys', {})[colname] = { + "column": fk.name, + "model": fk.expr.fullname + } + # Also add a selection of auto-generated attributes. sym = ctx.api.lookup_fully_qualified_or_none('sqlalchemy.sql.schema.Table') if sym: From c2686446d78600181367e643b3d656b58140dc6a Mon Sep 17 00:00:00 2001 From: Mehdi Date: Sun, 10 Feb 2019 17:32:36 +0100 Subject: [PATCH 2/7] better describing of table --- sqlmypy.py | 47 +++++++++++++++++++++++++++++++++-------------- 1 file changed, 33 insertions(+), 14 deletions(-) diff --git a/sqlmypy.py b/sqlmypy.py index 621efe1..7ed9bd3 100644 --- a/sqlmypy.py +++ b/sqlmypy.py @@ -115,26 +115,45 @@ def add_model_init_hook(ctx: ClassDefContext) -> None: if not (isinstance(stmt, AssignmentStmt) and len(stmt.lvalues) == 1 and isinstance(stmt.lvalues[0], NameExpr)): continue + # We currently only handle setting __tablename__ as a class attribute, and not through a property. if stmt.lvalues[0].name == "__tablename__" and isinstance(stmt.rvalue, StrExpr): ctx.cls.info.metadata.setdefault('sqlalchemy', {})['tablename'] = stmt.rvalue.value if isinstance(stmt.rvalue, CallExpr) and stmt.rvalue.callee.fullname == COLUMN_NAME: - colname = stmt.lvalues[0].name - has_explicit_colname = stmt.rvalue - ctx.cls.info.metadata.setdefault('sqlalchemy', {}).setdefault('columns', []).append(colname) + # Save columns. The name of a column on the db side can be different from the one inside the SA model. + sa_column_name = stmt.lvalues[0].name + + db_column_name = None # type: Optional[str] + if 'name' in stmt.rvalue.arg_names: + name_str_expr = stmt.rvalue.args[stmt.rvalue.arg_names.index('name')] + assert isinstance(name_str_expr, StrExpr) + db_column_name = name_str_expr.value + else: + if len(stmt.rvalue.args) >= 1 and isinstance(stmt.rvalue.args[0], StrExpr): + db_column_name = stmt.rvalue.args[0].value + + ctx.cls.info.metadata.setdefault('sqlalchemy', {}).setdefault('columns', []).append( + {"sa_name": sa_column_name, "db_name": db_column_name or sa_column_name} + ) + + # Save foreign keys. for arg in stmt.rvalue.args: if isinstance(arg, CallExpr) and arg.callee.fullname == FOREIGN_KEY_NAME and len(arg.args) >= 1: fk = arg.args[0] if isinstance(fk, StrExpr): - *_, parent_table, parent_col = fk.value.split(".") - ctx.cls.info.metadata.setdefault('sqlalchemy', {}).setdefault('foreign_keys', {})[colname] = { - "column": parent_col, - "table": parent_table + *r, parent_table_name, parent_db_col_name = fk.value.split(".") + assert len(r) <= 1 + ctx.cls.info.metadata.setdefault('sqlalchemy', {}).setdefault('foreign_keys', + {})[sa_column_name] = { + "db_name": parent_db_col_name, + "table_name": parent_table_name, + "schema": r[0] if r else None } elif isinstance(fk, MemberExpr): - ctx.cls.info.metadata.setdefault('sqlalchemy', {}).setdefault('foreign_keys', {})[colname] = { - "column": fk.name, - "model": fk.expr.fullname + ctx.cls.info.metadata.setdefault('sqlalchemy', {}).setdefault('foreign_keys', + {})[sa_column_name] = { + "sa_name": fk.name, + "model_fullname": fk.expr.fullname } # Also add a selection of auto-generated attributes. @@ -414,10 +433,10 @@ class User(Base): # We really need to add this to TypeChecker API def parse_bool(expr: Expression) -> Optional[bool]: if isinstance(expr, NameExpr): - if expr.fullname == 'builtins.True': - return True - if expr.fullname == 'builtins.False': - return False + if expr.fullname == 'builtins.True': + return True + if expr.fullname == 'builtins.False': + return False return None From ee7baf3c87c7559b4e3979d022d1c8eb2fd562b1 Mon Sep 17 00:00:00 2001 From: Mehdi Date: Sun, 10 Feb 2019 19:15:32 +0100 Subject: [PATCH 3/7] Try to guess if a relationship is iterable --- sqlmypy.py | 57 +++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/sqlmypy.py b/sqlmypy.py index 7ed9bd3..83728d0 100644 --- a/sqlmypy.py +++ b/sqlmypy.py @@ -117,7 +117,7 @@ def add_model_init_hook(ctx: ClassDefContext) -> None: # We currently only handle setting __tablename__ as a class attribute, and not through a property. if stmt.lvalues[0].name == "__tablename__" and isinstance(stmt.rvalue, StrExpr): - ctx.cls.info.metadata.setdefault('sqlalchemy', {})['tablename'] = stmt.rvalue.value + ctx.cls.info.metadata.setdefault('sqlalchemy', {})['table_name'] = stmt.rvalue.value if isinstance(stmt.rvalue, CallExpr) and stmt.rvalue.callee.fullname == COLUMN_NAME: # Save columns. The name of a column on the db side can be different from the one inside the SA model. @@ -363,6 +363,54 @@ def grouping_hook(ctx: FunctionContext) -> Type: return ctx.default_return_type +class IncompleteModelMetadata(Exception): + pass + + +def has_foreign_keys(local_model: TypeInfo, remote_model: TypeInfo) -> bool: + """Tells if `local_model` has a fk to `remote_model`. + Will raise an `IncompleteModelMetadata` if some mandatory metadata is missing. + """ + local_metadata = local_model.metadata.get("sqlalchemy", {}) + remote_metadata = remote_model.metadata.get("sqlalchemy", {}) + + for fk in local_metadata.get("foreign_keys", {}).values(): + if 'model_fullname' in fk and remote_model.fullname == fk['model_fullname']: + return True + if 'table_name' in fk: + if 'table_name' not in remote_metadata: + raise IncompleteModelMetadata + # TODO: handle different schemas + if remote_metadata['table_name'] == fk['table_name']: + return True + + return False + + +def is_relationship_iterable(ctx: FunctionContext, local_model: TypeInfo, remote_model: TypeInfo) -> bool: + """Tries to guess if the relationship is onetoone/onetomany/manytoone. + + Currently we handle the most current case, where a model relates to the other one through a relationship. + We also handle cases where secondaryjoin argument is provided. + We don't handle advanced usecases (foreign keys on both sides, primaryjoin, etc.). + """ + secondaryjoin = get_argument_by_name(ctx, 'secondaryjoin') + + if secondaryjoin is not None: + return True + + try: + can_be_many_to_one = has_foreign_keys(local_model, remote_model) + can_be_one_to_many = has_foreign_keys(remote_model, local_model) + + if not can_be_many_to_one and can_be_one_to_many: + return True + except IncompleteModelMetadata: + pass + + return False # Assume relationship is not iterable, if we weren't able to guess better. + + def relationship_hook(ctx: FunctionContext) -> Type: """Support basic use cases for relationships. @@ -415,10 +463,17 @@ class User(Base): # Something complex, stay silent for now. new_arg = AnyType(TypeOfAny.special_form) + current_model = ctx.api.scope.active_class() + assert current_model is not None + + # TODO: handle backref relationships + # We figured out, the model type. Now check if we need to wrap it in Iterable if uselist_arg: if parse_bool(uselist_arg): new_arg = ctx.api.named_generic_type('typing.Iterable', [new_arg]) + elif not isinstance(new_arg, AnyType) and is_relationship_iterable(ctx, current_model, new_arg.type): + new_arg = ctx.api.named_generic_type('typing.Iterable', [new_arg]) else: if has_annotation: # If there is an annotation we use it as a source of truth. From 25c3ef3ede5ce3e826428a960102469530ac9f32 Mon Sep 17 00:00:00 2001 From: Mehdi Date: Sun, 10 Feb 2019 19:28:15 +0100 Subject: [PATCH 4/7] mypy fixes --- sqlmypy.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/sqlmypy.py b/sqlmypy.py index 83728d0..e6d01ea 100644 --- a/sqlmypy.py +++ b/sqlmypy.py @@ -66,17 +66,17 @@ def get_function_hook(self, fullname: str) -> Optional[Callable[[FunctionContext return model_hook return None - def get_dynamic_class_hook(self, fullname: str): + def get_dynamic_class_hook(self, fullname: str) -> CB[DynamicClassDefContext]: if fullname == 'sqlalchemy.ext.declarative.api.declarative_base': return decl_info_hook return None - def get_class_decorator_hook(self, fullname: str): + def get_class_decorator_hook(self, fullname: str) -> CB[ClassDefContext]: if fullname == 'sqlalchemy.ext.declarative.api.as_declarative': return decl_deco_hook return None - def get_base_class_hook(self, fullname: str): + def get_base_class_hook(self, fullname: str) -> CB[ClassDefContext]: sym = self.lookup_fully_qualified(fullname) if sym and isinstance(sym.node, TypeInfo): if is_declarative(sym.node): @@ -119,7 +119,8 @@ def add_model_init_hook(ctx: ClassDefContext) -> None: if stmt.lvalues[0].name == "__tablename__" and isinstance(stmt.rvalue, StrExpr): ctx.cls.info.metadata.setdefault('sqlalchemy', {})['table_name'] = stmt.rvalue.value - if isinstance(stmt.rvalue, CallExpr) and stmt.rvalue.callee.fullname == COLUMN_NAME: + if (isinstance(stmt.rvalue, CallExpr) and isinstance(stmt.rvalue.callee, NameExpr) + and stmt.rvalue.callee.fullname == COLUMN_NAME): # Save columns. The name of a column on the db side can be different from the one inside the SA model. sa_column_name = stmt.lvalues[0].name @@ -138,7 +139,8 @@ def add_model_init_hook(ctx: ClassDefContext) -> None: # Save foreign keys. for arg in stmt.rvalue.args: - if isinstance(arg, CallExpr) and arg.callee.fullname == FOREIGN_KEY_NAME and len(arg.args) >= 1: + if (isinstance(arg, CallExpr) and isinstance(arg.callee, NameExpr) + and arg.callee.fullname == FOREIGN_KEY_NAME and len(arg.args) >= 1): fk = arg.args[0] if isinstance(fk, StrExpr): *r, parent_table_name, parent_db_col_name = fk.value.split(".") @@ -149,7 +151,7 @@ def add_model_init_hook(ctx: ClassDefContext) -> None: "table_name": parent_table_name, "schema": r[0] if r else None } - elif isinstance(fk, MemberExpr): + elif isinstance(fk, MemberExpr) and isinstance(fk.expr, NameExpr): ctx.cls.info.metadata.setdefault('sqlalchemy', {}).setdefault('foreign_keys', {})[sa_column_name] = { "sa_name": fk.name, @@ -463,7 +465,8 @@ class User(Base): # Something complex, stay silent for now. new_arg = AnyType(TypeOfAny.special_form) - current_model = ctx.api.scope.active_class() + # use private api + current_model = ctx.api.scope.active_class() # type: ignore # type: TypeInfo assert current_model is not None # TODO: handle backref relationships @@ -472,7 +475,7 @@ class User(Base): if uselist_arg: if parse_bool(uselist_arg): new_arg = ctx.api.named_generic_type('typing.Iterable', [new_arg]) - elif not isinstance(new_arg, AnyType) and is_relationship_iterable(ctx, current_model, new_arg.type): + elif isinstance(new_arg, Instance) and is_relationship_iterable(ctx, current_model, new_arg.type): new_arg = ctx.api.named_generic_type('typing.Iterable', [new_arg]) else: if has_annotation: From d5f2edec58497960cc7607dc8f489e3a3b12c29a Mon Sep 17 00:00:00 2001 From: Mehdi Date: Sun, 10 Feb 2019 19:43:45 +0100 Subject: [PATCH 5/7] fix --- sqlmypy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sqlmypy.py b/sqlmypy.py index e6d01ea..18cb4e3 100644 --- a/sqlmypy.py +++ b/sqlmypy.py @@ -377,7 +377,7 @@ def has_foreign_keys(local_model: TypeInfo, remote_model: TypeInfo) -> bool: remote_metadata = remote_model.metadata.get("sqlalchemy", {}) for fk in local_metadata.get("foreign_keys", {}).values(): - if 'model_fullname' in fk and remote_model.fullname == fk['model_fullname']: + if 'model_fullname' in fk and remote_model.fullname() == fk['model_fullname']: return True if 'table_name' in fk: if 'table_name' not in remote_metadata: From 605a997ca1dd8c0172ba77e1c1d3ed39f082a9ef Mon Sep 17 00:00:00 2001 From: Mehdi Date: Sun, 10 Feb 2019 19:43:53 +0100 Subject: [PATCH 6/7] add tests --- .../test-data/sqlalchemy-plugin-features.test | 90 +++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/test/test-data/sqlalchemy-plugin-features.test b/test/test-data/sqlalchemy-plugin-features.test index bcbbb3e..5f0b0f6 100644 --- a/test/test-data/sqlalchemy-plugin-features.test +++ b/test/test-data/sqlalchemy-plugin-features.test @@ -280,3 +280,93 @@ class M2(M1): Base = declarative_base(cls=(M1, M2)) # E: Not able to calculate MRO for declarative base reveal_type(Base) # E: Revealed type is 'Any' [out] + +[case testRelationshipIsGuessed] +from sqlalchemy import Column, Integer, String, ForeignKey +from sqlalchemy.orm import relationship +from sqlalchemy.ext.declarative import declarative_base + +Base = declarative_base() + +class Parent(Base): + __tablename__ = 'parents' + id = Column(Integer, primary_key=True) + name = Column(String) + + children = relationship("Child") + +class Child(Base): + __tablename__ = 'children' + id = Column(Integer, primary_key=True) + name = Column(String) + parent_id = Column(Integer, ForeignKey(Parent.id)) + + parent = relationship(Parent) + +child: Child +parent: Parent + +reveal_type(child.parent) # E: Revealed type is 'main.Parent*' +reveal_type(parent.children) # E: Revealed type is 'typing.Iterable*[main.Child]' + +[out] + +[case testRelationshipIsGuessed2] +from sqlalchemy import Column, Integer, String, ForeignKey +from sqlalchemy.orm import relationship +from sqlalchemy.ext.declarative import declarative_base + +Base = declarative_base() + +class Parent(Base): + __tablename__ = 'parents' + id = Column(Integer, primary_key=True) + name = Column(String) + + children = relationship("Child") + +class Child(Base): + __tablename__ = 'children' + id = Column(Integer, primary_key=True) + name = Column(String) + parent_id = Column(Integer, ForeignKey("parents.id")) + + parent = relationship(Parent) + +child: Child +parent: Parent + +reveal_type(child.parent) # E: Revealed type is 'main.Parent*' +reveal_type(parent.children) # E: Revealed type is 'typing.Iterable*[main.Child]' + +[out] + +[case testRelationshipIsGuessed3] +from sqlalchemy import Column, Integer, String, ForeignKey +from sqlalchemy.orm import relationship +from sqlalchemy.ext.declarative import declarative_base + +Base = declarative_base() + +class Parent(Base): + __tablename__ = 'parents' + id = Column(Integer, primary_key=True) + name = Column(String) + + children = relationship("Child") + +class Child(Base): + __tablename__ = 'children' + id = Column(Integer, primary_key=True) + name = Column(String) + parent_id = Column(Integer, ForeignKey("other_parents.id")) + + parent = relationship(Parent) + +child: Child +parent: Parent + +reveal_type(child.parent) # E: Revealed type is 'main.Parent*' +reveal_type(parent.children) # E: Revealed type is 'main.Child*' + +[out] \ No newline at end of file From 79d0504879e2aeeb5ff6241680997929101ce956 Mon Sep 17 00:00:00 2001 From: Mehdi Date: Sun, 10 Mar 2019 12:23:18 +0100 Subject: [PATCH 7/7] explain TODO --- sqlmypy.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sqlmypy.py b/sqlmypy.py index 18cb4e3..55a01b4 100644 --- a/sqlmypy.py +++ b/sqlmypy.py @@ -382,7 +382,8 @@ def has_foreign_keys(local_model: TypeInfo, remote_model: TypeInfo) -> bool: if 'table_name' in fk: if 'table_name' not in remote_metadata: raise IncompleteModelMetadata - # TODO: handle different schemas + # TODO: handle different schemas. + # It's not straightforward because schema can be specified in `__table_args__` or in metadata for example if remote_metadata['table_name'] == fk['table_name']: return True