From c678a4f69f6e5c1839303d66737180b18d91ecaf Mon Sep 17 00:00:00 2001 From: Rafal Stozek Date: Wed, 5 Jun 2019 20:05:32 +0200 Subject: [PATCH 1/4] Consider all columns in mro for __init__ method (#75) --- sqlmypy.py | 15 +++--- .../test-data/sqlalchemy-plugin-features.test | 47 +++++++++++++++++++ 2 files changed, 55 insertions(+), 7 deletions(-) diff --git a/sqlmypy.py b/sqlmypy.py index 9b66b7f..22fc2f1 100644 --- a/sqlmypy.py +++ b/sqlmypy.py @@ -199,7 +199,7 @@ def model_hook(ctx: FunctionContext) -> Type: argument types is 'Any'. """ assert isinstance(ctx.default_return_type, Instance) - model = ctx.default_return_type.type + model: TypeInfo = ctx.default_return_type.type metadata = model.metadata.get('sqlalchemy') if not metadata or not metadata.get('generated_init'): return ctx.default_return_type @@ -207,12 +207,13 @@ def model_hook(ctx: FunctionContext) -> Type: # Collect column names and types defined in the model # TODO: cache this? expected_types = {} # type: Dict[str, Type] - for name, sym in model.names.items(): - if isinstance(sym.node, Var) and isinstance(sym.node.type, Instance): - tp = sym.node.type - if tp.type.fullname() in (COLUMN_NAME, RELATIONSHIP_NAME): - assert len(tp.args) == 1 - expected_types[name] = tp.args[0] + for cls in model.mro[::-1]: + for name, sym in cls.names.items(): + if isinstance(sym.node, Var) and isinstance(sym.node.type, Instance): + tp = sym.node.type + if tp.type.fullname() in (COLUMN_NAME, RELATIONSHIP_NAME): + assert len(tp.args) == 1 + expected_types[name] = tp.args[0] assert len(ctx.arg_names) == 1 # only **kwargs in generated __init__ assert len(ctx.arg_types) == 1 diff --git a/test/test-data/sqlalchemy-plugin-features.test b/test/test-data/sqlalchemy-plugin-features.test index bcbbb3e..4b08df3 100644 --- a/test/test-data/sqlalchemy-plugin-features.test +++ b/test/test-data/sqlalchemy-plugin-features.test @@ -50,6 +50,53 @@ class Base: ... [out] +[case testModelInitMixin] + +from sqlalchemy import Column, Integer, String, DateTime +from sqlalchemy.ext.declarative import declarative_base + +Base = declarative_base() + +class HasId: + id = Column(Integer, primary_key=True) + + +class User(Base, HasId): + __tablename__ = 'users' + + name = Column(String, nullable=False) + + +user = User(id=123, name="John Doe") + +[out] + +[case testModelInitProperMro] + +from sqlalchemy import Column, Integer, String, DateTime +from sqlalchemy.ext.declarative import declarative_base + +Base = declarative_base() + +class Defaults: + id = Column(Integer, primary_key=True) + + +class User(Base, Defaults): + __tablename__ = 'users' + + # By default mypy will complain about Column[str] not being compatible with Column[int]. + # Adding "ignore" should allow us to override column type. + id = Column(String, primary_key=True) # type: ignore + name = Column(String, nullable=False) + + +User(id='stringish-id') +User(id=123) + +[out] +main:21: error: Incompatible type for "id" of "User" (got "int", expected "str") + [case testModelInitRelationship] from typing import TYPE_CHECKING, List From 69da77a5d13f8bd408f817c3fa90089e9bf0044b Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Wed, 10 Jul 2019 16:03:31 +0100 Subject: [PATCH 2/4] Can't use type annotations We support Python 3.5 (also it is not needed here). --- sqlmypy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sqlmypy.py b/sqlmypy.py index 22fc2f1..4482389 100644 --- a/sqlmypy.py +++ b/sqlmypy.py @@ -199,7 +199,7 @@ def model_hook(ctx: FunctionContext) -> Type: argument types is 'Any'. """ assert isinstance(ctx.default_return_type, Instance) - model: TypeInfo = ctx.default_return_type.type + model = ctx.default_return_type.type metadata = model.metadata.get('sqlalchemy') if not metadata or not metadata.get('generated_init'): return ctx.default_return_type From 7f0bb632c81abf4eff7c65cc0c96e7ce8cedef37 Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Wed, 10 Jul 2019 16:07:19 +0100 Subject: [PATCH 3/4] Use less space in tests Also small updates --- test/test-data/sqlalchemy-plugin-features.test | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/test/test-data/sqlalchemy-plugin-features.test b/test/test-data/sqlalchemy-plugin-features.test index 4b08df3..e11b078 100644 --- a/test/test-data/sqlalchemy-plugin-features.test +++ b/test/test-data/sqlalchemy-plugin-features.test @@ -60,15 +60,13 @@ Base = declarative_base() class HasId: id = Column(Integer, primary_key=True) - class User(Base, HasId): __tablename__ = 'users' name = Column(String, nullable=False) - user = User(id=123, name="John Doe") - +reveal_type(user.id) # N: Revealed type is 'builtins.int' [out] [case testModelInitProperMro] @@ -81,7 +79,6 @@ Base = declarative_base() class Defaults: id = Column(Integer, primary_key=True) - class User(Base, Defaults): __tablename__ = 'users' @@ -90,12 +87,9 @@ class User(Base, Defaults): id = Column(String, primary_key=True) # type: ignore name = Column(String, nullable=False) - User(id='stringish-id') -User(id=123) - +User(id=123) # E: Incompatible type for "id" of "User" (got "int", expected "str") [out] -main:21: error: Incompatible type for "id" of "User" (got "int", expected "str") [case testModelInitRelationship] from typing import TYPE_CHECKING, List From ad19ff70ca3623eafab8240fad624932e2ec075f Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Wed, 10 Jul 2019 16:37:16 +0100 Subject: [PATCH 4/4] Update test output --- test/test-data/sqlalchemy-plugin-features.test | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test-data/sqlalchemy-plugin-features.test b/test/test-data/sqlalchemy-plugin-features.test index e11b078..3eaab8d 100644 --- a/test/test-data/sqlalchemy-plugin-features.test +++ b/test/test-data/sqlalchemy-plugin-features.test @@ -66,7 +66,7 @@ class User(Base, HasId): name = Column(String, nullable=False) user = User(id=123, name="John Doe") -reveal_type(user.id) # N: Revealed type is 'builtins.int' +reveal_type(user.id) # N: Revealed type is 'builtins.int*' [out] [case testModelInitProperMro]