Lazy import of numpy/torch etc. (#253)

* Rewrite extensions in serializable to load torch/numpy lazily

* Add file with test requirements

* Add a dumb test for save functions

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Add numpy and pytorch as test dependencies

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Fix small typing issues and bug with save fn

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Remove duplicate tests from test_yaml.py

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

---------

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>
Co-authored-by: Fabrice Normandin <normandf@mila.quebec>

Author: breuleux

requirements-test.txt

@@ -0,0 +1,3 @@
+matplotlib
+numpy
+orion

setup.py

@@ -20,6 +20,8 @@
         "pytest",
         "pytest-xdist",
         "pytest-regressions",
+        "numpy",
+        "torch",
     ],
     "yaml": ["pyyaml"],
 }

simple_parsing/helpers/serialization/serializable.py

@@ -13,6 +13,8 @@
 from types import ModuleType
 from typing import IO, Any, Callable, ClassVar, TypeVar, Union
 
+from typing_extensions import Protocol
+
 from simple_parsing.utils import (
     DataclassT,
     all_subclasses,
@@ -57,6 +59,97 @@ def ordered_dict_representer(dumper: yaml.Dumper, instance: OrderedDict) -> yaml
     pass
 
 
+class FormatExtension(Protocol):
+    binary: ClassVar[bool] = False
+
+    @staticmethod
+    def load(fp: IO) -> Any:
+        ...
+
+    @staticmethod
+    def dump(obj: Any, io: IO) -> None:
+        ...
+
+
+class JSONExtension(FormatExtension):
+    load = staticmethod(json.load)
+    dump = staticmethod(json.dump)
+
+
+class PickleExtension(FormatExtension):
+    binary: ClassVar[bool] = True
+    load: ClassVar[Callable[[IO], Any]] = staticmethod(pickle.load)
+    dump: ClassVar[Callable[[Any, IO[bytes]], None]] = staticmethod(pickle.dump)
+
+
+class YamlExtension(FormatExtension):
+    def load(self, io: IO) -> Any:
+        import yaml
+
+        return yaml.safe_load(io)
+
+    def dump(self, obj: Any, io: IO, **kwargs) -> None:
+        import yaml
+
+        return yaml.dump(obj, io, **kwargs)
+
+
+class NumpyExtension(FormatExtension):
+    binary: bool = True
+
+    def load(self, io: IO) -> Any:
+        import numpy
+
+        obj = numpy.load(io, allow_pickle=True)
+        if isinstance(obj, numpy.ndarray) and obj.dtype == object:
+            obj = obj.item()
+        return obj
+
+    def dump(self, obj: Any, io: IO[bytes], **kwargs) -> None:
+        import numpy
+
+        return numpy.save(io, obj, **kwargs)
+
+
+class TorchExtension(FormatExtension):
+    binary: bool = True
+
+    def load(self, io: IO) -> None:
+        import torch  # type: ignore
+
+        return torch.load(io)
+
+    def dump(self, obj: Any, io: IO, **kwargs) -> None:
+        import torch  # type: ignore
+
+        return torch.save(obj, io, **kwargs)
+
+
+json_extension = JSONExtension()
+yaml_extension = YamlExtension()
+
+
+extensions: dict[str, FormatExtension] = {
+    ".json": JSONExtension(),
+    ".pkl": PickleExtension(),
+    ".yaml": YamlExtension(),
+    ".yml": YamlExtension(),
+    ".npy": NumpyExtension(),
+    ".pth": TorchExtension(),
+}
+
+
+def get_extension(path: str | Path) -> FormatExtension:
+    path = Path(path)
+    if path.suffix in extensions:
+        return extensions[path.suffix]
+    else:
+        raise RuntimeError(
+            f"Cannot load to/save from a {path.suffix} file because "
+            "this extension is not registered in the extensions dictionary."
+        )
+
+
 class SerializableMixin:
     """Makes a dataclass serializable to and from dictionaries.
 
@@ -248,17 +341,17 @@ def load_yaml(
         """
         return load_yaml(cls, path, load_fn=load_fn, drop_extra_fields=drop_extra_fields, **kwargs)
 
-    def save(self, path: str | Path, dump_fn=None) -> None:
-        save(self, path=path, dump_fn=dump_fn)
+    def save(self, path: str | Path, format: FormatExtension | None = None) -> None:
+        save(self, path=path, format=format)
 
-    def _save(self, path: str | Path, dump_fn: DumpFn = json.dump, **kwargs) -> None:
-        save(self, path=path, dump_fn=partial(dump_fn, **kwargs))
+    def _save(self, path: str | Path, format: FormatExtension = json_extension, **kwargs) -> None:
+        save(self, path=path, format=format, **kwargs)
 
     def save_yaml(self, path: str | Path, dump_fn: DumpFn | None = None, **kwargs) -> None:
-        save_yaml(self, path, dump_fn=dump_fn, **kwargs)
+        save_yaml(self, path, **kwargs)
 
-    def save_json(self, path: str | Path, dump_fn=json.dump, **kwargs) -> None:
-        save_json(self, path, dump_fn=dump_fn, **kwargs)
+    def save_json(self, path: str | Path, **kwargs) -> None:
+        save_json(self, path, **kwargs)
 
     @classmethod
     def loads(
@@ -484,51 +577,6 @@ def loads_yaml(
     return loads(cls, s, drop_extra_fields=drop_extra_fields, load_fn=partial(load_fn, **kwargs))
 
 
-extensions_to_loading_fn: dict[str, Callable[[IO], Any]] = {
-    ".json": json.load,
-    ".pkl": pickle.load,
-}
-extensions_to_read_mode: dict[str, str] = {".pkl": "rb"}
-
-extensions_to_write_mode: dict[str, str] = {".pkl": "wb"}
-extensions_to_dump_fn: dict[str, Callable[[Any, IO], None]] = {
-    ".json": json.dump,
-    ".pkl": pickle.dump,
-}
-try:
-    import yaml
-
-    extensions_to_loading_fn[".yaml"] = yaml.safe_load
-    extensions_to_loading_fn[".yml"] = yaml.safe_load
-    extensions_to_dump_fn[".yaml"] = yaml.dump
-    extensions_to_dump_fn[".yml"] = yaml.dump
-
-
-except ImportError:
-    pass
-
-try:
-    import numpy  # type: ignore
-
-    extensions_to_loading_fn[".npy"] = numpy.load
-    extensions_to_dump_fn[".npy"] = numpy.save
-    extensions_to_read_mode[".npy"] = "rb"
-    extensions_to_write_mode[".npy"] = "wb"
-
-except ImportError:
-    pass
-
-try:
-    import torch  # type: ignore
-
-    extensions_to_loading_fn[".pth"] = torch.load
-    extensions_to_dump_fn[".pth"] = torch.save
-    extensions_to_read_mode[".pth"] = "rb"
-    extensions_to_write_mode[".pth"] = "wb"
-except ImportError:
-    pass
-
-
 def read_file(path: str | Path) -> dict:
     """Returns the contents of the given file as a dictionary.
     Uses the right function depending on `path.suffix`:
@@ -540,66 +588,33 @@ def read_file(path: str | Path) -> dict:
         ".pkl": pickle.load,
     }
     """
-    path = Path(path)
-    if path.suffix in extensions_to_loading_fn:
-        load_fn = extensions_to_loading_fn[path.suffix]
-    else:
-        raise RuntimeError(
-            f"Unable to determine what function to use in order to load "
-            f"path {path} into a dictionary since the path's extension isn't registered in the "
-            f"`extensions_to_loading_fn` dictionary..."
-        )
-    mode = extensions_to_read_mode.get(path.suffix, "r")
-    with open(path, mode=mode) as f:
-        return load_fn(f)
+    format = get_extension(path)
+    with open(path, mode="rb" if format.binary else "r") as f:
+        return format.load(f)
 
 
 def save(
     obj: Any,
     path: str | Path,
-    dump_fn: Callable[[dict, IO], None] | None = None,
+    format: FormatExtension | None = None,
     save_dc_types: bool = False,
+    **kwargs,
 ) -> None:
-    """Save the given dataclass or dictionary to the given file.
-
-    Note: The `encode` function is applied to all the object fields to get serializable values,
-    like so:
-    - obj -> encode -> "raw" values (dicts, strings, ints, etc) -> `dump_fn` ([json/yaml/etc].dumps) -> string
-    """
-    path = Path(path)
-
+    """Save the given dataclass or dictionary to the given file."""
     if not isinstance(obj, dict):
         obj = to_dict(obj, save_dc_types=save_dc_types)
+    if format is None:
+        format = get_extension(path)
+    with open(path, mode="wb" if format.binary else "w") as f:
+        return format.dump(obj, f, **kwargs)
 
-    if dump_fn:
-        save_fn = dump_fn
-    elif path.suffix in extensions_to_dump_fn:
-        save_fn = extensions_to_dump_fn[path.suffix]
-    else:
-        raise RuntimeError(
-            f"Unable to determine what function to use in order to save obj {obj} to path {path},"
-            f"since the path's extension isn't registered in the "
-            f"`extensions_to_dump_fn` dictionary..."
-        )
-    mode = extensions_to_write_mode.get(path.suffix, "w")
-    with open(path, mode=mode) as f:
-        return save_fn(obj, f)
 
+def save_yaml(obj, path: str | Path, **kwargs) -> None:
+    save(obj, path, format=yaml_extension, **kwargs)
 
-def save_yaml(
-    obj, path: str | Path, dump_fn: DumpFn | None = None, save_dc_types: bool = False, **kwargs
-) -> None:
-    import yaml
 
-    if dump_fn is None:
-        dump_fn = yaml.dump
-    save(obj, path, dump_fn=partial(dump_fn, **kwargs), save_dc_types=save_dc_types)
-
-
-def save_json(
-    obj, path: str | Path, dump_fn: DumpFn = json.dump, save_dc_types: bool = False, **kwargs
-) -> None:
-    save(obj, path, dump_fn=partial(dump_fn, **kwargs), save_dc_types=save_dc_types)
+def save_json(obj, path: str | Path, **kwargs) -> None:
+    save(obj, path, format=json_extension, **kwargs)
 
 
 def load_yaml(

test/helpers/test_save.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+
+from ..nesting.example_use_cases import HyperParameters
+
+
+def test_save_yaml(tmpdir: Path):
+    hparams = HyperParameters.setup("")
+    tmp_path = Path(tmpdir / "temp.yml")
+    hparams.save_yaml(tmp_path)
+
+    _hparams = HyperParameters.load_yaml(tmp_path)
+    assert hparams == _hparams
+
+
+def test_save_json(tmpdir: Path):
+    hparams = HyperParameters.setup("")
+    tmp_path = Path(tmpdir / "temp.json")
+    hparams.save_yaml(tmp_path)
+    _hparams = HyperParameters.load_yaml(tmp_path)
+    assert hparams == _hparams
+
+
+def test_save_yml(tmpdir: Path):
+    hparams = HyperParameters.setup("")
+    tmp_path = Path(tmpdir / "temp.yml")
+    hparams.save(tmp_path)
+
+    _hparams = HyperParameters.load(tmp_path)
+    assert hparams == _hparams
+
+
+def test_save_pickle(tmpdir: Path):
+    hparams = HyperParameters.setup("")
+    tmp_path = Path(tmpdir / "temp.pkl")
+    hparams.save(tmp_path)
+
+    _hparams = HyperParameters.load(tmp_path)
+    assert hparams == _hparams
+
+
+def test_save_numpy(tmpdir: Path):
+    hparams = HyperParameters.setup("")
+    tmp_path = Path(tmpdir / "temp.npy")
+    hparams.save(tmp_path)
+
+    _hparams = HyperParameters.load(tmp_path)
+    assert hparams == _hparams
+
+
+def test_save_torch(tmpdir: Path):
+    hparams = HyperParameters.setup("")
+    tmp_path = Path(tmpdir / "temp.pth")
+    hparams.save(tmp_path)
+
+    _hparams = HyperParameters.load(tmp_path)
+    assert hparams == _hparams

test/utils/test_yaml.py

@@ -1,14 +1,11 @@
 """ Tests for serialization to/from yaml files. """
 import textwrap
 from dataclasses import dataclass
-from pathlib import Path
 from typing import List
 
 from simple_parsing import list_field
 from simple_parsing.helpers.serialization import YamlSerializable
 
-from ..nesting.example_use_cases import HyperParameters
-
 
 @dataclass
 class Point(YamlSerializable):
@@ -65,32 +62,6 @@ def test_dumps_loads():
     )
 
 
-def test_save_yaml(tmpdir: Path):
-    hparams = HyperParameters.setup("")
-    tmp_path = Path(tmpdir / "temp.yml")
-    hparams.save_yaml(tmp_path)
-
-    _hparams = HyperParameters.load_yaml(tmp_path)
-    assert hparams == _hparams
-
-
-def test_save_json(tmpdir: Path):
-    hparams = HyperParameters.setup("")
-    tmp_path = Path(tmpdir / "temp.json")
-    hparams.save_yaml(tmp_path)
-    _hparams = HyperParameters.load_yaml(tmp_path)
-    assert hparams == _hparams
-
-
-def test_save_yml(tmpdir: Path):
-    hparams = HyperParameters.setup("")
-    tmp_path = Path(tmpdir / "temp.yml")
-    hparams.save(tmp_path)
-
-    _hparams = HyperParameters.load(tmp_path)
-    assert hparams == _hparams
-
-
 # def test_save_yml(HyperParameters, tmpdir: Path):
 #     hparams = HyperParameters.setup("")
 #     tmp_path = Path(tmpdir / "temp.pth")