Skip to content

Conversation

KunxiSun
Copy link

@KunxiSun KunxiSun commented Apr 3, 2025

I add a IntEnum type decorator which I used in my private project.

I wish it helps!

from enum import IntEnum
class HeroStatus(IntEnum):
    ACTIVE = 1
    DISABLE = 2    

from sqlmodel import IntEnum
class Hero(SQLModel):
    hero_status: HeroStatus = Field(sa_type=IntEnum(HeroStatus))


# Save as Integer in database, but load as IntEnum in program
user.hero_status == HeroStatus.ACTIVE      

@KunxiSun

This comment was marked as resolved.

@svlandeg svlandeg changed the title Add IntEnum for sqltypes ✨ Add IntEnum for sqltypes Apr 4, 2025
@svlandeg svlandeg added the feature New feature or request label Apr 4, 2025
@KunxiSun

This comment was marked as resolved.

@svlandeg

This comment was marked as resolved.

@KunxiSun
Copy link
Author

@KunxiSun: Yes - those are unrelated to your PR. We're looking into it here.

Thanks to let me know

Copy link
Contributor

@YuriiMotov YuriiMotov left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@KunxiSun, thanks for your interest and efforts!

I think this feature is quite useful.
There seems to be no simple solution to configure int enums this way without creating custom TypeDecorator. This PR will provide such solution.

Please take a look at my in-code comments.
Also, we need to update these tests to check statements for new field

Comment on lines +40 to +44
if value is None:
return None

result = self.enum_type(value)
return result
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if value is None:
return None
result = self.enum_type(value)
return result
return None if (value is None) else self.enum_type(value)

Comment on lines +51 to +55
if value is None:
return None

result = value.value
return result
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if value is None:
return None
result = value.value
return result
return None if (value is None) else value.value

@@ -14,3 +17,39 @@ def load_dialect_impl(self, dialect: Dialect) -> "types.TypeEngine[Any]":
if impl.length is None and dialect.name == "mysql":
return dialect.type_descriptor(types.String(self.mysql_default_length))
return super().load_dialect_impl(dialect)


class IntEnum(types.TypeDecorator): # type: ignore
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
class IntEnum(types.TypeDecorator): # type: ignore
class IntEnum(types.TypeDecorator[Optional[_TIntEnum]]):


self.enum_type = enum_type

def process_result_value( # type: ignore[override]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def process_result_value( # type: ignore[override]
def process_result_value(

@@ -1,8 +1,11 @@
from typing import Any, cast
from enum import IntEnum as _IntEnum
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
from enum import IntEnum as _IntEnum
import enum


from sqlalchemy import types
from sqlalchemy.engine.interfaces import Dialect

_TIntEnum = TypeVar("_TIntEnum", bound="_IntEnum")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
_TIntEnum = TypeVar("_TIntEnum", bound="_IntEnum")
_TIntEnum = TypeVar("_TIntEnum", bound=enum.IntEnum)

super().__init__(*args, **kwargs)

# validate the input enum type
if not issubclass(enum_type, _IntEnum):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if not issubclass(enum_type, _IntEnum):
if not issubclass(enum_type, enum.IntEnum):

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
feature New feature or request waiting
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants