Merge branch 'feat/IntEnum' of https://github.com/KunxiSun/sqlmodel into feat/IntEnum

Author: KunxiSun

sqlmodel/sql/sqltypes.py

@@ -19,6 +19,7 @@ def load_dialect_impl(self, dialect: Dialect) -> "types.TypeEngine[Any]":
             return dialect.type_descriptor(types.String(self.mysql_default_length))
         return super().load_dialect_impl(dialect)
 
+
 class IntEnum(types.TypeDecorator):  # type: ignore
     """TypeDecorator for Integer-enum conversion.
 
@@ -30,7 +31,7 @@ class IntEnum(types.TypeDecorator):  # type: ignore
     Example:
         >>> class HeroStatus(enum.IntEnum):
         ...     ACTIVE = 1
-        ...     DISABLE = 2    
+        ...     DISABLE = 2
         >>>>
         >>> from sqlmodel import IntEnum
         >>> class Hero(SQLModel):
@@ -51,9 +52,7 @@ def __init__(self, enum_type: _TIntEnum, *args, **kwargs):
 
         # validate the input enum type
         if not issubclass(enum_type, _IntEnum):
-            raise TypeError(
-                f"Input must be enum.IntEnum"
-            )
+            raise TypeError("Input must be enum.IntEnum")
 
         self.enum_type = enum_type
 

tests/test_enums.py

@@ -78,7 +78,7 @@ def test_json_schema_flat_model_pydantic_v1():
                 "description": "An enumeration.",
                 "enum": [1, 3],
                 "type": "int",
-            }
+            },
         },
     }
 
@@ -106,7 +106,7 @@ def test_json_schema_inherit_model_pydantic_v1():
                 "description": "An int enumeration.",
                 "enum": [1, 3],
                 "type": "int",
-            }
+            },
         },
     }
 

tests/test_enums_models.py

@@ -1,7 +1,7 @@
 import enum
 import uuid
 
-from sqlmodel import Field, SQLModel, IntEnum
+from sqlmodel import Field, IntEnum, SQLModel
 
 
 class MyEnum1(str, enum.Enum):
@@ -13,15 +13,18 @@ class MyEnum2(str, enum.Enum):
     C = "C"
     D = "D"
 
+
 class MyEnum3(enum.IntEnum):
     E = 1
     F = 2
 
+
 class BaseModel(SQLModel):
     id: uuid.UUID = Field(primary_key=True)
     enum_field: MyEnum2
     int_enum_field: MyEnum3
 
+
 class FlatModel(SQLModel, table=True):
     id: uuid.UUID = Field(primary_key=True)
     enum_field: MyEnum1