Fix enum type resolution

miquido · Sep 5, 2024 · d150cc1 · d150cc1
1 parent 3e053bd
commit d150cc1
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "draive"
 description = "Framework designed to simplify and accelerate the development of LLM-based applications."
-version = "0.28.4"
+version = "0.28.5"
 readme = "README.md"
 maintainers = [
  { name = "Kacper Kaliński", email = "[email protected]" },

diff --git a/src/draive/parameters/specification.py b/src/draive/parameters/specification.py
@@ -456,13 +456,13 @@ def parameter_specification( # noqa: C901, PLR0912, PLR0915, PLR0911, PLR0913
  raise TypeError("Unsupported literal type annotation: %s", annotation)
 
  case enum_type if isinstance(enum_type, enum.EnumType):
- if isinstance(enum_type, enum.StrEnum):
+ if issubclass(enum_type, enum.StrEnum):
  specification = {
  "type": "string",
  "enum": list(enum_type),
  }
 
- elif isinstance(enum_type, enum.IntEnum):
+ elif issubclass(enum_type, enum.IntEnum):
  specification = {
  "type": "integer",
  "enum": list(enum_type),

diff --git a/src/draive/parameters/validation.py b/src/draive/parameters/validation.py
@@ -908,10 +908,10 @@ def parameter_validator[Value]( # noqa: C901, PLR0915, PLR0912, PLR0913
  validator = _prepare_any_enum_validator(resolved_args)
 
  case enum_type if isinstance(enum_type, enum.EnumType):
- if isinstance(enum_type, enum.StrEnum):
+ if issubclass(enum_type, enum.StrEnum):
  validator = _prepare_str_enum_validator(resolved_args)
 
- elif isinstance(enum_type, enum.IntEnum):
+ elif issubclass(enum_type, enum.IntEnum):
  validator = _prepare_int_enum_validator(resolved_args)
 
  else: