From 9e07c1c772b6526e98d981090f064273743db2e3 Mon Sep 17 00:00:00 2001
From: honglei <jhonglei@gmail.com>
Date: Tue, 15 Aug 2023 22:38:03 +0800
Subject: [PATCH] fix type check for sa_column

---
 sqlmodel/main.py         | 6 ++++--
 sqlmodel/sql/sqltypes.py | 1 -
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/sqlmodel/main.py b/sqlmodel/main.py
index fc99f00..ab1de9f 100644
--- a/sqlmodel/main.py
+++ b/sqlmodel/main.py
@@ -167,7 +167,7 @@ def Field(
     unique: bool = False,
     nullable: Union[bool, PydanticUndefinedType] = PydanticUndefined,
     index: Union[bool, PydanticUndefinedType] = PydanticUndefined,
-    sa_column: Union[Column, PydanticUndefinedType, types.FunctionType] = PydanticUndefined,  # type: ignore
+    sa_column: Union[Column, PydanticUndefinedType, Callable[[], Column]] = PydanticUndefined,  # type: ignore
     sa_column_args: Union[Sequence[Any], PydanticUndefinedType] = PydanticUndefined,
     sa_column_kwargs: Union[
         Mapping[str, Any], PydanticUndefinedType
@@ -525,7 +525,9 @@ def get_column_from_field(field: FieldInfo) -> Column:  # type: ignore
     if isinstance(sa_column, MappedColumn):
         return sa_column.column
     if isinstance(sa_column, types.FunctionType):
-        return sa_column()
+        col = sa_column()
+        assert isinstance(col, Column)
+        return col
     sa_type = get_sqlalchemy_type(field)
     primary_key = getattr(field, "primary_key", False)
     index = getattr(field, "index", PydanticUndefined)
diff --git a/sqlmodel/sql/sqltypes.py b/sqlmodel/sql/sqltypes.py
index da6551b..aa30950 100644
--- a/sqlmodel/sql/sqltypes.py
+++ b/sqlmodel/sql/sqltypes.py
@@ -8,7 +8,6 @@ from sqlalchemy.sql.type_api import TypeEngine
 
 
 class AutoString(types.TypeDecorator):  # type: ignore
-
     impl = types.String
     cache_ok = True
     mysql_default_length = 255