Skip to content

Commit 0a12128

Browse files
committed
Make Engine a str enum and normalize model_type
Update `Engine` to inherit from `str, Enum` so enum members behave like strings where needed. Also harden `from_model_type` by coercing non-string inputs (including enum-like values with `.value`) before lowercasing, and raise a clear `ValueError` when conversion is not possible.
1 parent 1d34a6a commit 0a12128

1 file changed

Lines changed: 7 additions & 1 deletion

File tree

dlclivegui/temp/engine.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
# or if we update dlclive.Engine to have these methods and use that instead of a separate enum here.
77
# The latter would be more cohesive but also creates a dependency from utils to dlclive,
88
# pending release of dlclive
9-
class Engine(Enum):
9+
class Engine(str, Enum):
1010
TENSORFLOW = "tensorflow"
1111
PYTORCH = "pytorch"
1212

@@ -26,6 +26,12 @@ def is_tensorflow_model_dir_path(model_path: str | Path) -> bool:
2626

2727
@classmethod
2828
def from_model_type(cls, model_type: str) -> "Engine":
29+
if not isinstance(model_type, str):
30+
try:
31+
model_type = getattr(model_type, "value", str(model_type))
32+
except Exception as e:
33+
raise ValueError(f"Could not convert model_type to string: {model_type}") from e
34+
2935
if model_type.lower() == "pytorch":
3036
return cls.PYTORCH
3137
elif model_type.lower() in ("tensorflow", "base", "tensorrt", "lite"):

0 commit comments

Comments
 (0)