Fix MetaData.bind typing (#43)

Author: bryanforbes

sqlalchemy-stubs/engine/__init__.pyi

@@ -1,3 +1,4 @@
+from typing import Any
 from . import default as default
 
 from .interfaces import (
@@ -7,7 +8,7 @@ from .interfaces import (
     ExecutionContext as ExecutionContext,
     ExceptionContext as ExceptionContext,
     Compiled as Compiled,
-    TypeCompiler as TypeCompiler
+    TypeCompiler as TypeCompiler,
 )
 
 from .base import (
@@ -29,5 +30,5 @@ from .result import (
     RowProxy as RowProxy,
 )
 
-def create_engine(*args, **kwargs): ...
-def engine_from_config(configuration, prefix: str = ..., **kwargs): ...
+def create_engine(*args: Any, **kwargs: Any) -> Engine: ...
+def engine_from_config(configuration: Any, prefix: str = ..., **kwargs: Any) -> Engine: ...

sqlalchemy-stubs/sql/schema.pyi

@@ -9,6 +9,7 @@ from .selectable import TableClause as TableClause
 from .type_api import TypeEngine
 from .. import util
 from ..engine import Engine, Connection, Connectable
+from ..engine.url import URL
 from .compiler import DDLCompiler
 import threading
 
@@ -298,18 +299,25 @@ class Index(DialectKWArgs, ColumnCollectionMixin, SchemaItem):
 
 DEFAULT_NAMING_CONVENTION: util.immutabledict[str, str] = ...
 
+class _MetaDataBind:
+    @overload
+    def __get__(self, instance: None, owner: Any) -> None: ...
+    @overload
+    def __get__(self, instance: MetaData, owner: Any) -> Optional[Union[Engine, Connection]]: ...
+    def __set__(self, instance: Any, value: Optional[Union[Engine, Connection, str, URL]]) -> None: ...
+
 class MetaData(SchemaItem):
     __visit_name__: str = ...
     tables: util.immutabledict[str, Table] = ...
     schema: Optional[str] = ...
     naming_convention: Mapping[Union[str, Index, Constraint], str] = ...
     info: Optional[Mapping[str, Any]] = ...
+    bind: _MetaDataBind = ...
     def __init__(self, bind: Optional[Union[Engine, Connection]] = ..., reflect: bool = ..., schema: Optional[str] = ...,
                  quote_schema: Optional[bool] = ..., naming_convention: Mapping[Union[str, Index, Constraint], str] = ...,
                  info: Optional[Mapping[str, Any]] = ...) -> None: ...
     def __contains__(self, table_or_key: Union[str, Table]) -> bool: ...
     def is_bound(self) -> bool: ...
-    def bind(self) -> Optional[Union[Engine, Connection]]: ...
     def clear(self) -> None: ...
     def remove(self, table: Table) -> None: ...
     @property
@@ -326,7 +334,7 @@ class MetaData(SchemaItem):
 class ThreadLocalMetaData(MetaData):
     __visit_name__: str = ...
     context: threading.local = ...
+    bind: _MetaDataBind = ...
     def __init__(self) -> None: ...
-    def bind(self) -> Optional[Union[Engine, Connection]]: ...
     def is_bound(self) -> bool: ...
     def dispose(self) -> None: ...

test/test-data/sqlalchemy-sql-schema.test

@@ -0,0 +1,16 @@
+[case testSchemaMetaData]
+from sqlalchemy import MetaData, create_engine
+from sqlalchemy.engine import Connection
+from sqlalchemy.engine.url import make_url
+
+m = MetaData()
+e = create_engine('postgresql://foo')
+c = Connection(e)
+m.bind = 'postgresql://foo'
+m.bind = make_url('postgresql://foo')
+reveal_type(m.bind)  # E: Revealed type is 'Union[sqlalchemy.engine.base.Engine, sqlalchemy.engine.base.Connection, None]'
+m.bind = e
+reveal_type(m.bind)  # E: Revealed type is 'sqlalchemy.engine.base.Engine'
+m.bind = c
+reveal_type(m.bind)  # E: Revealed type is 'sqlalchemy.engine.base.Connection'
+[out]

test/testsql.py

@@ -23,7 +23,8 @@ class SQLDataSuite(DataSuite):
     files = ['sqlalchemy-basics.test',
              'sqlalchemy-sql-elements.test',
              'sqlalchemy-sql-sqltypes.test',
-             'sqlalchemy-sql-selectable.test']
+             'sqlalchemy-sql-selectable.test',
+             'sqlalchemy-sql-schema.test']
     data_prefix = test_data_prefix
 
     def run_case(self, testcase: DataDrivenTestCase) -> None: