Add argument compressed to save/load functions

muspy/inputs/json.py

@@ -1,27 +1,47 @@
 """JSON input interface."""
+import gzip
 import json
 from pathlib import Path
-from typing import TextIO, Union
+from typing import Optional, TextIO, Union
 
 from ..music import Music
 
 
-def load_json(path: Union[str, Path, TextIO]) -> Music:
+def load_json(
+    path: Union[str, Path, TextIO], compressed: Optional[bool] = None
+) -> Music:
     """Load a JSON file into a Music object.
 
     Parameters
     ----------
     path : str, Path or TextIO
         Path to the file or the file to load.
+    compressed : bool, optional
+        Whether the file is a compressed JSON file (`.json.gz`). Has no
+        effect when `path` is a file object. Defaults to infer from the
+        extension (`.gz`).
 
     Returns
     -------
     :class:`muspy.Music`
         Loaded Music object.
 
+    Notes
+    -----
+    When a path is given, assume UTF-8 encoding and gzip compression if
+    `compressed=True`.
+
     """
     if isinstance(path, (str, Path)):
-        with open(str(path), encoding="utf-8") as f:
+        if compressed is None:
+            if str(path).lower().endswith(".gz"):
+                compressed = True
+            else:
+                compressed = False
+        if compressed:
+            with gzip.open(path, "rt", encoding="utf-8") as f:
+                return Music.from_dict(json.load(f))
+        with open(path, encoding="utf-8") as f:
             return Music.from_dict(json.load(f))
 
     return Music.from_dict(json.load(path))

muspy/inputs/wrappers.py

@@ -22,16 +22,23 @@
 from .yaml import load_yaml
 
 
-def load(path: Union[str, Path, TextIO], kind: Optional[str] = None) -> Music:
+def load(
+    path: Union[str, Path, TextIO], kind: Optional[str] = None, **kwargs
+) -> Music:
     """Load a JSON or a YAML file into a Music object.
 
+    This is a wrapper function for :func:`muspy.load_json` and
+    :func:`muspy.load_yaml`.
+
     Parameters
     ----------
     path : str, Path or TextIO
         Path to the file or the file to to load.
     kind : {'json', 'yaml'}, optional
-        Format to save (case-insensitive). Defaults to infer the format
-        from the extension.
+        Format to save. Defaults to infer from the extension.
+    **kwargs
+        Keyword arguments to pass to :func:`muspy.load_json` or
+        :func:`muspy.load_yaml`.
 
     Returns
     -------
@@ -40,25 +47,29 @@ def load(path: Union[str, Path, TextIO], kind: Optional[str] = None) -> Music:
 
     See Also
     --------
+    :func:`muspy.load_json` : Load a JSON file into a Music object.
+    :func:`muspy.load_yaml` : Load a YAML file into a Music object.
     :func:`muspy.read` :
         Read a MIDI/MusicXML/ABC file into a Music object.
 
     """
-    # pylint: disable=unused-argument
     if kind is None:
-        if str(path).lower().endswith(".json"):
+        if not isinstance(path, (str, Path)):
+            raise ValueError("Cannot infer file format from a file object.")
+        path_str = str(path).lower()
+        if path_str.endswith((".json", ".json.gz")):
             kind = "json"
-        elif str(path).lower().endswith((".yaml", ".yml")):
+        elif path_str.endswith((".yaml", ".yml", ".yaml.gz", ".yml.gz")):
             kind = "yaml"
         else:
             raise ValueError(
                 "Cannot infer file format from the extension (expect JSON or "
                 "YAML)."
             )
     if kind.lower() == "json":
-        return load_json(path)
+        return load_json(path, **kwargs)
     if kind.lower() == "yaml":
-        return load_yaml(path)
+        return load_yaml(path, **kwargs)
     raise ValueError(
         f"Expect `kind` to be 'json' or 'yaml', but got : {kind}."
     )
@@ -74,8 +85,7 @@ def read(
     path : str or Path
         Path to the file to read.
     kind : {'midi', 'musicxml', 'abc'}, optional
-        Format to save (case-insensitive). Defaults to infer the format
-        from the extension.
+        Format to save. Defaults to infer from the extension.
     **kwargs
         Keyword arguments to pass to :func:`muspy.read_midi`,
         :func:`muspy.read_musicxml` or :func:`read_abc`.
@@ -157,7 +167,7 @@ def from_representation(array: ndarray, kind: str, **kwargs) -> Music:
     array : :class:`numpy.ndarray`
         Array in a supported representation.
     kind : str, {'pitch', 'pianoroll', 'event', 'note'}
-        Data representation type (case-insensitive).
+        Data representation.
     **kwargs
         Keyword arguments to pass to
         :func:`muspy.from_pitch_representation`,

muspy/inputs/yaml.py

@@ -1,28 +1,48 @@
 """YAML input interface."""
+import gzip
 from pathlib import Path
-from typing import TextIO, Union
+from typing import Optional, TextIO, Union
 
 import yaml
 
 from ..music import Music
 
 
-def load_yaml(path: Union[str, Path, TextIO]) -> Music:
+def load_yaml(
+    path: Union[str, Path, TextIO], compressed: Optional[bool] = None
+) -> Music:
     """Load a YAML file into a Music object.
 
     Parameters
     ----------
     path : str, Path or TextIO
         Path to the file or the file to load.
+    compressed : bool, optional
+        Whether the file is a compressed YAML file (`.yaml.gz`). Has no
+        effect when `path` is a file object. Defaults to infer from the
+        extension (`.gz`).
 
     Returns
     -------
     :class:`muspy.Music`
         Loaded Music object.
 
+    Notes
+    -----
+    When a path is given, assume UTF-8 encoding and gzip compression if
+    `compressed=True`.
+
     """
     if isinstance(path, (str, Path)):
-        with open(str(path), encoding="utf-8") as f:
+        if compressed is None:
+            if str(path).lower().endswith(".gz"):
+                compressed = True
+            else:
+                compressed = False
+        if compressed:
+            with gzip.open(path, "rt", encoding="utf-8") as f:
+                return Music.from_dict(yaml.safe_load(f))
+        with open(path, encoding="utf-8") as f:
             return Music.from_dict(yaml.safe_load(f))
 
     return Music.from_dict(yaml.safe_load(path))

muspy/outputs/json.py

@@ -1,7 +1,8 @@
 """JSON output interface."""
+import gzip
 import json
 from pathlib import Path
-from typing import TYPE_CHECKING, TextIO, Union
+from typing import TYPE_CHECKING, Optional, TextIO, Union
 
 if TYPE_CHECKING:
     from ..music import Music
@@ -12,6 +13,7 @@ def save_json(
     music: "Music",
     skip_missing: bool = True,
     ensure_ascii: bool = False,
+    compressed: Optional[bool] = None,
     **kwargs
 ):
     """Save a Music object to a JSON file.
@@ -28,16 +30,34 @@ def save_json(
     ensure_ascii : bool
         Whether to escape non-ASCII characters. Will be passed to
         PyYAML's `yaml.dump`. Defaults to False.
+    compressed : bool, optional
+        Whether to save as a compressed JSON file (`.json.gz`). Has no
+        effect when `path` is a file object. Defaults to infer from the
+        extension (`.gz`).
     **kwargs
         Keyword arguments to pass to :py:func:`json.dumps`.
 
+    Notes
+    -----
+    When a path is given, use UTF-8 encoding and gzip compression if
+    `compressed=True`.
+
     """
     ordered_dict = music.to_ordered_dict(skip_missing=skip_missing)
     data = json.dumps(ordered_dict, ensure_ascii=ensure_ascii, **kwargs)
 
     if isinstance(path, (str, Path)):
-        with open(str(path), "w", encoding="utf-8") as f:
-            f.write(data)
+        if compressed is None:
+            if str(path).lower().endswith(".gz"):
+                compressed = True
+            else:
+                compressed = False
+        if compressed:
+            with gzip.open(path, "wt", encoding="utf-8") as f:
+                f.write(data)
+        else:
+            with open(path, "w", encoding="utf-8") as f:
+                f.write(data)
         return
 
     path.write(data)

muspy/outputs/wrappers.py

@@ -32,18 +32,25 @@ def save(
 ):
     """Save a Music object loselessly to a JSON or a YAML file.
 
+    This is a wrapper function for :func:`muspy.save_json` and
+    :func:`muspy.save_yaml`.
+
     Parameters
     ----------
     path : str, Path or TextIO
         Path or file to save the data.
     music : :class:`muspy.Music`
         Music object to save.
     kind : {'json', 'yaml'}, optional
-        Format to save (case-insensitive). Defaults to infer the format
-        from the extension.
+        Format to save. Defaults to infer from the extension.
+    **kwargs
+        Keyword arguments to pass to :func:`muspy.save_json` or
+        :func:`muspy.save_yaml`.
 
     See Also
     --------
+    :func:`muspy.save_json` : Save a Music object to a JSON file.
+    :func:`muspy.save_yaml` : Save a Music object to a YAML file.
     :func:`muspy.write` :
         Write a Music object to a MIDI/MusicXML/ABC/audio file.
 
@@ -57,9 +64,10 @@ def save(
     if kind is None:
         if not isinstance(path, (str, Path)):
             raise ValueError("Cannot infer file format from a file object.")
-        if str(path).lower().endswith(".json"):
+        path_str = str(path).lower()
+        if path_str.endswith((".json", ".json.gz")):
             kind = "json"
-        elif str(path).lower().endswith((".yaml", ".yml")):
+        elif path_str.endswith((".yaml", ".yml", ".yaml.gz", ".yml.gz")):
             kind = "yaml"
         else:
             raise ValueError(
@@ -90,8 +98,7 @@ def write(
     music : :class:`muspy.Music`
         Music object to convert.
     kind : {'midi', 'musicxml', 'abc', 'audio'}, optional
-        Format to save (case-insensitive). Defaults to infer the format
-        from the extension.
+        Format to save. Defaults to infer from the extension.
 
     See Also
     --------
@@ -142,7 +149,7 @@ def to_object(
     music : :class:`muspy.Music`
         Music object to convert.
     kind : str, {'music21', 'mido', 'pretty_midi', 'pypianoroll'}
-        Target class (case-insensitive).
+        Target class.
 
     Returns
     -------
@@ -174,7 +181,7 @@ def to_representation(music: "Music", kind: str, **kwargs) -> ndarray:
     music : :class:`muspy.Music`
         Music object to convert.
     kind : str, {'pitch', 'piano-roll', 'event', 'note'}
-        Target representation (case-insensitive).
+        Target representation.
 
     Returns
     -------

muspy/outputs/yaml.py

@@ -1,6 +1,7 @@
 """YAML output interface."""
+import gzip
 from pathlib import Path
-from typing import TYPE_CHECKING, TextIO, Union
+from typing import TYPE_CHECKING, Optional, TextIO, Union
 
 from ..utils import yaml_dump
 
@@ -13,6 +14,7 @@ def save_yaml(
     music: "Music",
     skip_missing: bool = True,
     allow_unicode: bool = True,
+    compressed: Optional[bool] = None,
     **kwargs
 ):
     """Save a Music object to a YAML file.
@@ -29,16 +31,34 @@ def save_yaml(
     allow_unicode : bool
         Whether to escape non-ASCII characters. Will be passed to
         :py:func:`json.dumps`. Defaults to False.
+    compressed : bool, optional
+        Whether to save as a compressed YAML file (`.yaml.gz`). Has no
+        effect when `path` is a file object. Defaults to infer from the
+        extension (`.gz`).
     **kwargs
-        Keyword arguments to pass to :py:func:`json.dumps`.
+        Keyword arguments to pass to `yaml.dump`.
+
+    Notes
+    -----
+    When a path is given, use UTF-8 encoding and gzip compression if
+    `compressed=True`.
 
     """
     ordered_dict = music.to_ordered_dict(skip_missing=skip_missing)
     data = yaml_dump(ordered_dict, allow_unicode=allow_unicode, **kwargs)
 
     if isinstance(path, (str, Path)):
-        with open(str(path), "w", encoding="utf-8") as f:
-            f.write(data)
+        if compressed is None:
+            if str(path).lower().endswith(".gz"):
+                compressed = True
+            else:
+                compressed = False
+        if compressed:
+            with gzip.open(path, "wt", encoding="utf-8") as f:
+                f.write(data)
+        else:
+            with open(path, "w", encoding="utf-8") as f:
+                f.write(data)
         return
 
     path.write(data)

tests/test_io.py

@@ -38,6 +38,14 @@ def test_save_json_path(tmp_path):
     check_music(loaded)
 
 
+def test_save_json_path_compressed(tmp_path):
+    music = muspy.load(TEST_JSON_PATH)
+    music.save(tmp_path / "test.json.gz")
+
+    loaded = muspy.load(tmp_path / "test.json.gz")
+    check_music(loaded)
+
+
 def test_save_json_file(tmp_path):
     music = muspy.load(TEST_JSON_PATH)
     with open(tmp_path / "test.json", "w", encoding="utf-8") as f:
@@ -55,6 +63,14 @@ def test_save_yaml_path(tmp_path):
     check_music(loaded)
 
 
+def test_save_yaml_path_compressed(tmp_path):
+    music = muspy.load(TEST_JSON_PATH)
+    music.save(tmp_path / "test.yaml.gz")
+
+    loaded = muspy.load(tmp_path / "test.yaml.gz")
+    check_music(loaded)
+
+
 def test_save_yaml_file(tmp_path):
     music = muspy.load(TEST_JSON_PATH)
     with open(tmp_path / "test.yaml", "w", encoding="utf-8") as f: