From 6ad07e26f8b8fc3da931ea572547bb6f2643e088 Mon Sep 17 00:00:00 2001
From: Federico Caselli <cfederico87@gmail.com>
Date: Tue, 16 May 2023 21:52:02 +0200
Subject: Define type for generic classes

Fixed typing use of :class:`~sqlalchemy.schema.Column` and other
generic SQLAlchemy classes.

Fixes: #1246
Change-Id: I5ee80395d626894a52e3395c9986213289576355
---
 tools/write_pyi.py | 6 ++++++
 1 file changed, 6 insertions(+)

(limited to 'tools')

diff --git a/tools/write_pyi.py b/tools/write_pyi.py
index 82ceead..da5b484 100644
--- a/tools/write_pyi.py
+++ b/tools/write_pyi.py
@@ -41,6 +41,7 @@ TRIM_MODULE = [
     "sqlalchemy.sql.type_api.",
     "sqlalchemy.sql.functions.",
     "sqlalchemy.sql.dml.",
+    "typing."
 ]
 ADDITIONAL_ENV = {
     "MigrationContext": MigrationContext,
@@ -180,6 +181,11 @@ def _generate_stub_for_meth(
             retval = annotation.__qualname__
         elif isinstance(annotation, typing.TypeVar):
             retval = annotation.__name__
+        elif hasattr(annotation, "__args__") and hasattr(
+            annotation, "__origin__"
+        ):
+            # generic class
+            retval = str(annotation)
         else:
             retval = annotation
 
-- 
cgit v1.2.1