Add support for labeling Literal types correctly.

matthewwardrop · Aug 15, 2024 · ee96b15 · ee96b15
1 parent fe0ea6f
commit ee96b15
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/spec_classes/utils/type_checking.py b/spec_classes/utils/type_checking.py
@@ -149,6 +149,10 @@ def type_label(attr_type: Type) -> str:
     """
     if hasattr(attr_type, "__origin__"):  # Generics
         label = type_label(attr_type.__origin__)
+        if attr_type.__origin__ is Literal:
+            return (
+                f"{label}[{', '.join(repr(arg) for arg in attr_type.__args__)}]"
+            )
         if hasattr(attr_type, "__args__") and not any(
             isinstance(arg, TypeVar) for arg in attr_type.__args__
         ):

diff --git a/tests/utils/test_type_checking.py b/tests/utils/test_type_checking.py
@@ -122,6 +122,7 @@ def test_attr_type_label(self):
         assert type_label(KeyedList[KeyedSpec, str]()) == "KeyedList[KeyedSpec, str]"
         assert type_label(Any) == "Any"
         assert type_label(Union[Any, Dict[str, int]]) == "Union[Any, dict[str, int]]"
+        assert type_label(Literal["a", True, 3, b"c"]) == "Literal['a', True, 3, b'c']"
 
     def test_type_instantiate(self):
         assert type_instantiate(str) == ""