From 21512f479c43e3ca0abb5a59a88d1e75171c707e Mon Sep 17 00:00:00 2001
From: Viicos <65306057+Viicos@users.noreply.github.com>
Date: Sat, 25 Mar 2023 22:02:22 -0400
Subject: Fix type annotation for `url`

<!-- Provide a general summary of your proposed changes in the Title field above -->
Left the `pyi` generated file untouched.
### Description
<!-- Describe your changes in detail -->

### Checklist
<!-- go over following points. check them with an `x` if they do apply, (they turn into clickable checkboxes once the PR is submitted, so no need to do everything at once)

-->

This pull request is:

- [x] A documentation / typographical error fix
	- Good to go, no issue or tests are needed
- [ ] A short code fix
	- please include the issue number, and create an issue if none exists, which
	  must include a complete example of the issue.  one line code fixes without an
	  issue and demonstration will not be accepted.
	- Please include: `Fixes: #<issue number>` in the commit message
	- please include tests.   one line code fixes without tests will not be accepted.
- [ ] A new feature implementation
	- please include the issue number, and create an issue if none exists, which must
	  include a complete example of how the feature would look.
	- Please include: `Fixes: #<issue number>` in the commit message
	- please include tests.

**Have a nice day!**

Closes: #1209
Pull-request: https://github.com/sqlalchemy/alembic/pull/1209
Pull-request-sha: d15cb6e28d5bb934d139ced87e713a9d77505346

Change-Id: I8548806a693b150f1417a102e6498119524008a0
---
 alembic/context.pyi            | 3 ++-
 alembic/runtime/environment.py | 3 ++-
 alembic/runtime/migration.py   | 3 ++-
 tools/write_pyi.py             | 1 +
 4 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/alembic/context.pyi b/alembic/context.pyi
index 142a0c7..a262638 100644
--- a/alembic/context.pyi
+++ b/alembic/context.pyi
@@ -17,6 +17,7 @@ from typing import Union
 
 if TYPE_CHECKING:
     from sqlalchemy.engine.base import Connection
+    from sqlalchemy.engine.url import URL
     from sqlalchemy.sql.elements import ClauseElement
     from sqlalchemy.sql.schema import MetaData
 
@@ -75,7 +76,7 @@ config: Config
 
 def configure(
     connection: Optional[Connection] = None,
-    url: Optional[str] = None,
+    url: Union[str, URL, None] = None,
     dialect_name: Optional[str] = None,
     dialect_opts: Optional[Dict[str, Any]] = None,
     transactional_ddl: Optional[bool] = None,
diff --git a/alembic/runtime/environment.py b/alembic/runtime/environment.py
index 0f9d3a5..c2fa11a 100644
--- a/alembic/runtime/environment.py
+++ b/alembic/runtime/environment.py
@@ -20,6 +20,7 @@ from ..operations import Operations
 if TYPE_CHECKING:
     from typing import Literal
 
+    from sqlalchemy.engine import URL
     from sqlalchemy.engine.base import Connection
     from sqlalchemy.sql.elements import ClauseElement
     from sqlalchemy.sql.schema import MetaData
@@ -334,7 +335,7 @@ class EnvironmentContext(util.ModuleClsProxy):
     def configure(
         self,
         connection: Optional[Connection] = None,
-        url: Optional[str] = None,
+        url: Optional[Union[str, URL]] = None,
         dialect_name: Optional[str] = None,
         dialect_opts: Optional[Dict[str, Any]] = None,
         transactional_ddl: Optional[bool] = None,
diff --git a/alembic/runtime/migration.py b/alembic/runtime/migration.py
index a615651..4e2d062 100644
--- a/alembic/runtime/migration.py
+++ b/alembic/runtime/migration.py
@@ -34,6 +34,7 @@ from ..util.compat import EncodedIO
 
 if TYPE_CHECKING:
     from sqlalchemy.engine import Dialect
+    from sqlalchemy.engine import URL
     from sqlalchemy.engine.base import Connection
     from sqlalchemy.engine.base import Transaction
     from sqlalchemy.engine.mock import MockConnection
@@ -217,7 +218,7 @@ class MigrationContext:
     def configure(
         cls,
         connection: Optional[Connection] = None,
-        url: Optional[str] = None,
+        url: Optional[Union[str, URL]] = None,
         dialect_name: Optional[str] = None,
         dialect: Optional[Dialect] = None,
         environment_context: Optional[EnvironmentContext] = None,
diff --git a/tools/write_pyi.py b/tools/write_pyi.py
index aec1813..fcf9d36 100644
--- a/tools/write_pyi.py
+++ b/tools/write_pyi.py
@@ -33,6 +33,7 @@ TRIM_MODULE = [
     "alembic.operations.base.",
     "alembic.operations.ops.",
     "sqlalchemy.engine.base.",
+    "sqlalchemy.engine.url.",
     "sqlalchemy.sql.schema.",
     "sqlalchemy.sql.selectable.",
     "sqlalchemy.sql.elements.",
-- 
cgit v1.2.1