Define type for generic classes

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

Fixes: #1246
Change-Id: I5ee80395d626894a52e3395c9986213289576355

1 files changed, 6 insertions, 0 deletions

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