Fix ModelHubMixin coders (#2291)

* update HubMixinTest with union and optional custom type

* enable ModelHubMixin to handle union and optional custom type

* add docstring for _is_optional_type helper function

* Refactor helper functions and add independent unit tests. Refactor decode_arg

* Restrict UnionType check to python3.10 and above. Minor style updates.

* Only branch to pipe tests when version >= python3.10

* Change pipe operator tests to str + eval

---------

Co-authored-by: Lucain <[email protected]>

Author: gorold

src/huggingface_hub/hub_mixin.py

@@ -4,7 +4,19 @@
 import warnings
 from dataclasses import asdict, dataclass, is_dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union, get_args
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    get_args,
+)
 
 from .constants import CONFIG_NAME, PYTORCH_WEIGHTS_NAME, SAFETENSORS_SINGLE_FILE
 from .file_download import hf_hub_download
@@ -16,8 +28,10 @@
     SoftTemporaryDirectory,
     is_jsonable,
     is_safetensors_available,
+    is_simple_optional_type,
     is_torch_available,
     logging,
+    unwrap_simple_optional_type,
     validate_hf_hub_args,
 )
 
@@ -336,14 +350,20 @@ def _encode_arg(cls, arg: Any) -> Any:
         """Encode an argument into a JSON serializable format."""
         for type_, (encoder, _) in cls._hub_mixin_coders.items():
             if isinstance(arg, type_):
+                if arg is None:
+                    return None
                 return encoder(arg)
         return arg
 
     @classmethod
-    def _decode_arg(cls, expected_type: Type[ARGS_T], value: Any) -> ARGS_T:
+    def _decode_arg(cls, expected_type: Type[ARGS_T], value: Any) -> Optional[ARGS_T]:
         """Decode a JSON serializable value into an argument."""
+        if is_simple_optional_type(expected_type):
+            if value is None:
+                return None
+            expected_type = unwrap_simple_optional_type(expected_type)
         for type_, (_, decoder) in cls._hub_mixin_coders.items():
-            if issubclass(expected_type, type_):
+            if inspect.isclass(expected_type) and issubclass(expected_type, type_):
                 return decoder(value)
         return value
 

src/huggingface_hub/utils/__init__.py

@@ -113,7 +113,7 @@
 from ._subprocess import capture_output, run_interactive_subprocess, run_subprocess
 from ._telemetry import send_telemetry
 from ._token import get_token
-from ._typing import is_jsonable
+from ._typing import is_jsonable, is_simple_optional_type, unwrap_simple_optional_type
 from ._validators import (
     smoothly_deprecate_use_auth_token,
     validate_hf_hub_args,

src/huggingface_hub/utils/_typing.py

@@ -14,7 +14,15 @@
 # limitations under the License.
 """Handle typing imports based on system compatibility."""
 
-from typing import Any, Callable, Literal, TypeVar
+import sys
+from typing import Any, Callable, List, Literal, Type, TypeVar, Union, get_args, get_origin
+
+
+UNION_TYPES: List[Any] = [Union]
+if sys.version_info >= (3, 10):
+    from types import UnionType
+
+    UNION_TYPES += [UnionType]
 
 
 HTTP_METHOD_T = Literal["GET", "OPTIONS", "HEAD", "POST", "PUT", "PATCH", "DELETE"]
@@ -48,3 +56,20 @@ def is_jsonable(obj: Any) -> bool:
         return False
     except RecursionError:
         return False
+
+
+def is_simple_optional_type(type_: Type) -> bool:
+    """Check if a type is optional, i.e. Optional[Type] or Union[Type, None] or Type | None, where Type is a non-composite type."""
+    if get_origin(type_) in UNION_TYPES:
+        union_args = get_args(type_)
+        if len(union_args) == 2 and type(None) in union_args:
+            return True
+    return False
+
+
+def unwrap_simple_optional_type(optional_type: Type) -> Type:
+    """Unwraps a simple optional type, i.e. returns Type from Optional[Type]."""
+    for arg in get_args(optional_type):
+        if arg is not type(None):
+            return arg
+    raise ValueError(f"'{optional_type}' is not an optional type")

tests/test_hub_mixin.py

@@ -144,11 +144,22 @@ class DummyModelWithCustomTypes(
     },
 ):
     def __init__(
-        self, foo: int, bar: str, custom: CustomType, custom_default: CustomType = CustomType("default"), **kwargs
+        self,
+        foo: int,
+        bar: str,
+        baz: Union[int, str],
+        custom: CustomType,
+        optional_custom_1: Optional[CustomType],
+        optional_custom_2: Optional[CustomType],
+        custom_default: CustomType = CustomType("default"),
+        **kwargs,
     ):
         self.foo = foo
         self.bar = bar
+        self.baz = baz
         self.custom = custom
+        self.optional_custom_1 = optional_custom_1
+        self.optional_custom_2 = optional_custom_2
         self.custom_default = custom_default
 
     @classmethod
@@ -406,21 +417,34 @@ def test_from_pretrained_when_cls_is_a_dataclass(self):
         assert not hasattr(model, "other")
 
     def test_from_cls_with_custom_type(self):
-        model = DummyModelWithCustomTypes(1, bar="bar", custom=CustomType("custom"))
+        model = DummyModelWithCustomTypes(
+            1,
+            bar="bar",
+            baz=1.0,
+            custom=CustomType("custom"),
+            optional_custom_1=CustomType("optional"),
+            optional_custom_2=None,
+        )
         model.save_pretrained(self.cache_dir)
 
         config = json.loads((self.cache_dir / "config.json").read_text())
         assert config == {
             "foo": 1,
             "bar": "bar",
+            "baz": 1.0,
             "custom": {"value": "custom"},
+            "optional_custom_1": {"value": "optional"},
+            "optional_custom_2": None,
             "custom_default": {"value": "default"},
         }
 
         model_reloaded = DummyModelWithCustomTypes.from_pretrained(self.cache_dir)
         assert model_reloaded.foo == 1
         assert model_reloaded.bar == "bar"
+        assert model_reloaded.baz == 1.0
         assert model_reloaded.custom.value == "custom"
+        assert model_reloaded.optional_custom_1 is not None and model_reloaded.optional_custom_1.value == "optional"
+        assert model_reloaded.optional_custom_2 is None
         assert model_reloaded.custom_default.value == "default"
 
     def test_inherited_class(self):

tests/test_utils_typing.py

@@ -1,14 +1,20 @@
 import json
+import sys
+from typing import Optional, Type, Union
 
 import pytest
 
-from huggingface_hub.utils._typing import is_jsonable
+from huggingface_hub.utils._typing import is_jsonable, is_simple_optional_type, unwrap_simple_optional_type
 
 
 class NotSerializableClass:
     pass
 
 
+class CustomType:
+    pass
+
+
 OBJ_WITH_CIRCULAR_REF = {"hello": "world"}
 OBJ_WITH_CIRCULAR_REF["recursive"] = OBJ_WITH_CIRCULAR_REF
 
@@ -47,3 +53,76 @@ def test_is_jsonable_failure(data):
     assert not is_jsonable(data)
     with pytest.raises((TypeError, ValueError)):
         json.dumps(data)
+
+
+@pytest.mark.parametrize(
+    "type_, is_optional",
+    [
+        (Optional[int], True),
+        (Union[None, int], True),
+        (Union[int, None], True),
+        (Optional[CustomType], True),
+        (Union[None, CustomType], True),
+        (Union[CustomType, None], True),
+        (int, False),
+        (None, False),
+        (Union[int, float, None], False),
+        (Union[Union[int, float], None], False),
+        (Optional[Union[int, float]], False),
+    ],
+)
+def test_is_simple_optional_type(type_: Type, is_optional: bool):
+    assert is_simple_optional_type(type_) is is_optional
+
+
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires python3.10 or higher")
+@pytest.mark.parametrize(
+    "type_, is_optional",
+    [
+        ("int | None", True),
+        ("None | int", True),
+        ("CustomType | None", True),
+        ("None | CustomType", True),
+        ("int | float", False),
+        ("int | float | None", False),
+        ("(int | float) | None", False),
+        ("Union[int, float] | None", False),
+    ],
+)
+def test_is_simple_optional_type_pipe(type_: str, is_optional: bool):
+    assert is_simple_optional_type(eval(type_)) is is_optional
+
+
+@pytest.mark.parametrize(
+    "optional_type, inner_type",
+    [
+        (Optional[int], int),
+        (Union[int, None], int),
+        (Union[None, int], int),
+        (Optional[CustomType], CustomType),
+        (Union[CustomType, None], CustomType),
+        (Union[None, CustomType], CustomType),
+    ],
+)
+def test_unwrap_simple_optional_type(optional_type: Type, inner_type: Type):
+    assert unwrap_simple_optional_type(optional_type) is inner_type
+
+
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires python3.10 or higher")
+@pytest.mark.parametrize(
+    "optional_type, inner_type",
+    [
+        ("None | int", int),
+        ("int | None", int),
+        ("None | CustomType", CustomType),
+        ("CustomType | None", CustomType),
+    ],
+)
+def test_unwrap_simple_optional_type_pipe(optional_type: str, inner_type: Type):
+    assert unwrap_simple_optional_type(eval(optional_type)) is inner_type
+
+
+@pytest.mark.parametrize("non_optional_type", [int, None, CustomType])
+def test_unwrap_simple_optional_type_fail(non_optional_type: Type):
+    with pytest.raises(ValueError):
+        unwrap_simple_optional_type(non_optional_type)