Merge branch 'bugfix-57/literal-satype' of https://github.com/nsaccente/sqlmodel into bugfix-57/literal-satype

Author: svlandeg

sqlmodel/_compat.py

@@ -219,6 +219,13 @@ def get_sa_type_from_type_annotation(annotation: Any) -> Any:
             # Optional unions are allowed
             use_type = bases[0] if bases[0] is not NoneType else bases[1]
             return get_sa_type_from_type_annotation(use_type)
+        if origin is Literal:
+            literal_args = get_args(annotation)
+            if all(isinstance(arg, bool) for arg in literal_args):  # all bools
+                return bool
+            if all(isinstance(arg, int) for arg in literal_args):  # all ints
+                return int
+            return str
         return origin
 
     def get_sa_type_from_field(field: Any) -> Any:
@@ -471,7 +478,13 @@ def is_field_noneable(field: "FieldInfo") -> bool:
 
     def get_sa_type_from_field(field: Any) -> Any:
         if get_origin(field.type_) is Literal:
-            return Literal
+            literal_args = get_args(field.type_)
+            if all(isinstance(arg, bool) for arg in literal_args):  # all bools
+                return bool
+            if all(isinstance(arg, int) for arg in literal_args):  # all ints
+                return int
+            return str
+
         if isinstance(field.type_, type) and field.shape == SHAPE_SINGLETON:
             return field.type_
         raise ValueError(f"The field {field.name} has no matching SQLAlchemy type")

sqlmodel/main.py

@@ -664,9 +664,6 @@ def get_sqlalchemy_type(field: Any) -> Any:
     type_ = get_sa_type_from_field(field)
     metadata = get_field_metadata(field)
 
-    # Checks for `Literal` type annotation
-    if type_ is Literal:
-        return AutoString
     # Check enums first as an enum can also be a str, needed by Pydantic/FastAPI
     if issubclass(type_, Enum):
         return sa_Enum(type_)

tests/test_main.py

@@ -128,23 +128,48 @@ class Hero(SQLModel, table=True):
         assert hero_rusty_man.team.name == "Preventers"
 
 
-def test_literal_typehints_are_treated_as_strings(clear_sqlmodel):
+def test_literal_str(clear_sqlmodel, caplog):
     """Test https://github.com/fastapi/sqlmodel/issues/57"""
 
-    class Hero(SQLModel, table=True):
+    class Model(SQLModel, table=True):
         id: Optional[int] = Field(default=None, primary_key=True)
-        name: str = Field(unique=True)
-        weakness: Literal["Kryptonite", "Dehydration", "Munchies"]
-
-    superguy = Hero(name="Superguy", weakness="Kryptonite")
+        all_str: Literal["a", "b", "c"]
+        mixed: Literal["yes", "no", 1, 0]
+        all_int: Literal[1, 2, 3]
+        int_bool: Literal[0, 1, True, False]
+        all_bool: Literal[True, False]
+
+    obj = Model(
+        all_str="a",
+        mixed="yes",
+        all_int=1,
+        int_bool=True,
+        all_bool=False,
+    )
 
     engine = create_engine("sqlite://", echo=True)
 
     SQLModel.metadata.create_all(engine)
 
+    # Check DDL
+    assert "all_str VARCHAR NOT NULL" in caplog.text
+    assert "mixed VARCHAR NOT NULL" in caplog.text
+    assert "all_int INTEGER NOT NULL" in caplog.text
+    assert "int_bool INTEGER NOT NULL" in caplog.text
+    assert "all_bool BOOLEAN NOT NULL" in caplog.text
+
+    # Check query
     with Session(engine) as session:
-        session.add(superguy)
+        session.add(obj)
         session.commit()
-        session.refresh(superguy)
-        assert superguy.weakness == "Kryptonite"
-        assert isinstance(superguy.weakness, str)
+        session.refresh(obj)
+        assert isinstance(obj.all_str, str)
+        assert obj.all_str == "a"
+        assert isinstance(obj.mixed, str)
+        assert obj.mixed == "yes"
+        assert isinstance(obj.all_int, int)
+        assert obj.all_int == 1
+        assert isinstance(obj.int_bool, int)
+        assert obj.int_bool == 1
+        assert isinstance(obj.all_bool, bool)
+        assert obj.all_bool is False