Update docstring to suggest raw string (#1186)

* Update docstring to suggest raw string

When escaping the `:` for `op.execute` if a normal string is used, you need to escape the backslash, eg. `'\\:colon_value'` or using a raw string.
For the docs I feel it is nicer to show the raw string.

* Update stub documentation

Change-Id: Ia605c6c036fe82cebff9b427333404f1c59ea74d

---------

Co-authored-by: CaselIT <cfederico87@gmail.com>

1 files changed, 3 insertions, 0 deletions

diff --git a/tools/write_pyi.py b/tools/write_pyi.py
index e3feb36..376163b 100644
--- a/tools/write_pyi.py
+++ b/tools/write_pyi.py
@@ -13,6 +13,7 @@ sys.path.append(str(Path(__file__).parent.parent))
 if True:  # avoid flake/zimports messing with the order
     from alembic.operations.base import Operations
     from alembic.runtime.environment import EnvironmentContext
+    from alembic.runtime.migration import MigrationContext
     from alembic.script.write_hooks import console_scripts
     from alembic.util.compat import inspect_formatargspec
     from alembic.util.compat import inspect_getfullargspec
@@ -40,6 +41,7 @@ TRIM_MODULE = [
     "sqlalchemy.sql.dml.",
 ]
 CONTEXT_MANAGERS = {"op": ["batch_alter_table"]}
+ADDITIONAL_ENV = {"MigrationContext": MigrationContext}
 
 
 def generate_pyi_for_proxy(
@@ -92,6 +94,7 @@ def generate_pyi_for_proxy(
             **sa.sql.schema.__dict__,
             **sa.__dict__,
             **sa.types.__dict__,
+            **ADDITIONAL_ENV,
             **ops.__dict__,
             **module.__dict__,
         }