Improve engine.mock (#101)

Author: bryanforbes

sqlalchemy-stubs/engine/mock.pyi

@@ -1,22 +1,32 @@
 from typing import Any
+from typing import NoReturn
+from typing import TypeVar
 
-from . import base as base
-from .. import util as util
-from ..sql import ddl as ddl
+from . import base
+from . import interfaces
+from ..sql.compiler import Compiled
+
+_TMockConnection = TypeVar("_TMockConnection", bound=MockConnection)
 
 class MockConnection(base.Connectable):
-    def __init__(self, dialect: Any, execute: Any) -> None: ...
-    engine: Any = ...
-    dialect: Any = ...
-    name: Any = ...
-    def schema_for_object(self, obj: Any): ...
-    def connect(self, **kwargs: Any): ...
-    def execution_options(self, **kw: Any): ...
-    def compiler(self, statement: Any, parameters: Any, **kwargs: Any): ...
+    def __init__(self, dialect: interfaces.Dialect, execute: Any) -> None: ...
+    engine: base.Engine = ...
+    dialect: interfaces.Dialect = ...
+    name: str = ...
+    def schema_for_object(self, obj: Any) -> str: ...
+    def connect(self: _TMockConnection, **kwargs: Any) -> _TMockConnection: ...  # type: ignore[override]
+    def execution_options(
+        self: _TMockConnection, **opt: Any
+    ) -> _TMockConnection: ...
+    def compiler(
+        self, statement: Any, parameters: Any, **kwargs: Any
+    ) -> Compiled: ...
     def create(self, entity: Any, **kwargs: Any) -> None: ...
     def drop(self, entity: Any, **kwargs: Any) -> None: ...
     def execute(
         self, object_: Any, *multiparams: Any, **params: Any
-    ) -> None: ...
+    ) -> NoReturn: ...
 
-def create_mock_engine(url: Any, executor: Any, **kw: Any): ...
+def create_mock_engine(
+    url: Any, executor: Any, **kw: Any
+) -> MockConnection: ...