Added missing attributes from context stubs.

Fixes: #900
Change-Id: I2e21b6f363bf6405bd472c9aca04fbbc3c9cc384

diff --git a/alembic/context.pyi b/alembic/context.pyi
index c70f5fcc7c32989048823eea3e74cbbba93f7ea2..5c29d3aef0dd350df732a4bf795c1e0e5e5c3470 100644 (file)
--- a/alembic/context.pyi
+++ b/alembic/context.pyi
@@ -13,8 +13,10 @@ if TYPE_CHECKING:
     from sqlalchemy.engine.base import Connection
     from sqlalchemy.sql.schema import MetaData
 
+    from .config import Config
     from .runtime.migration import _ProxyTransaction
     from .runtime.migration import MigrationContext
+    from .script import ScriptDirectory
 
 ### end imports ###
 
@@ -62,6 +64,8 @@ def begin_transaction() -> Union["_ProxyTransaction", ContextManager]:
 
     """
 
+config: Config
+
 def configure(
     connection: Optional["Connection"] = None,
     url: Optional[str] = None,
@@ -715,6 +719,8 @@ def run_migrations(**kw) -> None:
 
     """
 
+script: ScriptDirectory
+
 def static_output(text):
     """Emit text directly to the "offline" SQL stream.
 

diff --git a/docs/build/unreleased/900.rst b/docs/build/unreleased/900.rst
new file mode 100644 (file)
index 0000000..88f9f17
--- /dev/null
+++ b/docs/build/unreleased/900.rst
@@ -0,0 +1,5 @@
+.. change::
+    :tags: bug, typing
+    :tickets: 900
+
+    Added missing attributes from context stubs.

diff --git a/tools/write_pyi.py b/tools/write_pyi.py
index df2a55fd056079ae6a022b1472463fdb2359a2aa..60728a8d72d6fd90576bfbd32d0b6690da616940 100644 (file)
--- a/tools/write_pyi.py
+++ b/tools/write_pyi.py
@@ -19,11 +19,9 @@ if True:  # avoid flake/zimports messing with the order
 IGNORE_ITEMS = {
     "op": {"context", "create_module_class_proxy"},
     "context": {
-        "config",
         "create_module_class_proxy",
         "get_impl",
         "requires_connection",
-        "script",
     },
 }
 
@@ -95,7 +93,8 @@ def generate_pyi_for_proxy(
 
 
 def _generate_stub_for_attr(cls, name, printer):
-    printer.writeline(f"{name}: Any")
+    type_ = cls.__annotations__.get(name, "Any")
+    printer.writeline(f"{name}: {type_}")
 
 
 def _generate_stub_for_meth(cls, name, printer):