FEAT CLI conversion

setup.py

@@ -7,7 +7,7 @@
 # This is a bit (!) hackish: we are setting a global variable so that the
 # main modelcard __init__ can detect if it is being loaded by the setup
 # routine, to avoid attempting to load components.
-builtins.__SKOPS_SETUP__ = True
+builtins.__SKOPS_SETUP__ = True  # type: ignore
 
 
 import skops  # noqa
@@ -34,6 +34,14 @@
 
 
 def setup_package():
+    package_data = dict(
+        entry_points={
+            "console_scripts": [
+                "skops = skops.utils.entrypoint:main_entrypoint",
+            ],
+        }
+    )
+
     metadata = dict(
         name=DISTNAME,
         maintainer=MAINTAINER,
@@ -73,7 +81,7 @@ def setup_package():
         include_package_data=True,
     )
 
-    setup(**metadata)
+    setup(**package_data, **metadata)
 
 
 if __name__ == "__main__":

skops/io/_cli.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+import pathlib
+import pickle as pkl
+from typing import Optional
+
+from skops.io import dumps, get_untrusted_types
+
+
+def _convert(
+    input_file: os.PathLike, output_dir: pathlib.Path, is_trusted: bool = False
+):
+    model_name = pathlib.Path(input_file).stem
+
+    logging.info(f"Converting {model_name}")
+
+    with open(input_file, "rb") as f:
+        obj = pkl.load(f)
+    skops_dump = dumps(obj)
+
+    untrusted_types = get_untrusted_types(data=skops_dump)
+
+    if not untrusted_types:
+        logging.debug(f"No unsafe types found in {model_name}.")
+    else:
+        untrusted_str = "\n".join(untrusted_types)
+
+        logging.warning(
+            "Unsafe Types Detected!\n"
+            f"While converting {model_name}, "
+            "the following unsafe types were found: \n"
+            f"{untrusted_str}\n"
+        )
+
+        if not is_trusted:
+            logging.warning(
+                f"Model {model_name} will not be converted due to unsafe types.\n"
+                "To convert this anyway, add `-t` to this command."
+            )
+            return
+
+    with open(output_dir / f"{model_name}.skops", "wb") as out_file:
+        logging.info(f"Writing to {output_dir / model_name}.skops")
+        out_file.write(skops_dump)
+
+
+def main_convert(
+    command_line_args: Optional[list[str]] = None,
+    parent: Optional[argparse.ArgumentParser] = None,
+):
+    parents = [parent] if parent else []
+
+    parser = argparse.ArgumentParser(
+        description="Convert input Pickle files to .skops files", parents=parents
+    )
+
+    parser.add_argument("inputs", nargs="+", help="Input files to convert.")
+    parser.add_argument(
+        "-t",
+        "--trusted",
+        help=(
+            "Automatically trust all files, "
+            "and convert all inputs, "
+            "even if an untrusted type is detected."
+        ),
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "-d",
+        "--debug",
+        help="Enable debug logging.",
+        action="store_const",
+        dest="loglevel",
+        const=logging.DEBUG,
+        default=logging.WARNING,
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        help="Enable verbose logging.",
+        action="store_const",
+        dest="loglevel",
+        const=logging.INFO,
+    )
+    parser.add_argument(
+        "--output-dir",
+        help=(
+            "Specify a directory to save converted files to. "
+            "Default will save files to the current working directory."
+        ),
+        type=str,
+        default=None,
+    )
+    args = parser.parse_args(command_line_args)
+
+    output_dir = args.output_dir
+    if not output_dir:
+        output_dir = pathlib.Path.cwd()
+    else:
+        output_dir = pathlib.Path(output_dir)
+
+    logging.basicConfig(format="%(message)s", level=args.loglevel)
+
+    for input_file in args.inputs:
+        _convert(
+            input_file=input_file,
+            output_dir=output_dir,
+            is_trusted=args.trusted,
+        )

skops/io/tests/test_cli.py

@@ -0,0 +1,120 @@
+import logging
+import os.path
+import pathlib
+import pickle as pkl
+from typing import Optional
+from unittest import mock
+
+import numpy as np
+import pytest
+
+import skops
+from skops.io._cli import _convert
+
+
+class MockUnsafeType:
+    def __init__(self):
+        pass
+
+
+class TestConvert:
+    model_name = "some_model_name"
+
+    @pytest.fixture
+    def safe_obj(self):
+        return np.ndarray([1, 2, 3, 4])
+
+    @pytest.fixture
+    def unsafe_obj(self):
+        return MockUnsafeType()
+
+    @pytest.fixture
+    def pkl_path(self, tmp_path):
+        return tmp_path / f"{self.model_name}.pkl"
+
+    @pytest.fixture
+    def skops_path(self, tmp_path):
+        return tmp_path / f"{self.model_name}.skops"
+
+    @pytest.fixture
+    def write_safe_file(self, pkl_path, safe_obj):
+        with open(pkl_path, "wb") as f:
+            pkl.dump(safe_obj, f)
+
+    @pytest.fixture
+    def write_unsafe_file(self, pkl_path, unsafe_obj):
+        with open(pkl_path, "wb") as f:
+            pkl.dump(unsafe_obj, f)
+
+    def test_base_case_works_as_expected(
+        self, pkl_path, tmp_path, skops_path, write_safe_file, safe_obj
+    ):
+        _convert(pkl_path, tmp_path)
+        persisted_obj = skops.io.load(skops_path)
+        assert np.array_equal(persisted_obj, safe_obj)
+
+    def test_unsafe_case_works_as_expected(
+        self, pkl_path, tmp_path, skops_path, write_unsafe_file, caplog
+    ):
+        caplog.set_level(logging.WARNING)
+        _convert(pkl_path, tmp_path)
+        assert not os.path.isfile(skops_path)
+
+        # check logging has warned that an unsafe type was found
+        assert MockUnsafeType.__name__ in caplog.text
+
+
+class TestMainConvert:
+    @staticmethod
+    def assert_called_correctly(
+        mock_convert: mock.MagicMock,
+        paths: list,
+        output_dir: Optional[pathlib.Path] = pathlib.Path.cwd(),
+        trusted: Optional[bool] = False,
+    ):
+        assert mock_convert.call_count == len(paths)
+
+        mock_convert.assert_has_calls(
+            [
+                mock.call(input_file=p, output_dir=output_dir, is_trusted=trusted)
+                for p in paths
+            ]
+        )
+
+    @mock.patch("skops.io._cli._convert")
+    def test_base_works_as_expected(self, mock_convert: mock.MagicMock):
+        args = [
+            "123.pkl",
+            "abc.pkl",
+        ]
+
+        skops.io._cli.main_convert(command_line_args=args)
+        self.assert_called_correctly(mock_convert, args)
+
+    @mock.patch("skops.io._cli._convert")
+    @pytest.mark.parametrize("trusted_flag", ["-t", "--trusted"])
+    def test_with_trusted_works_as_expected(
+        self, mock_convert: mock.MagicMock, trusted_flag
+    ):
+        paths = ["abc.123", "234.567", "d/object_1.pkl"]
+        args = paths + [trusted_flag]
+        skops.io._cli.main_convert(command_line_args=args)
+        self.assert_called_correctly(mock_convert, paths=paths, trusted=True)
+
+    @mock.patch("skops.io._cli._convert")
+    @pytest.mark.parametrize(
+        "output_dir, expected_dir",
+        [("a/b/c", pathlib.Path("a/b/c")), (None, pathlib.Path.cwd())],
+    )
+    def test_with_output_dir_works_as_expected(
+        self, mock_convert: mock.MagicMock, output_dir, expected_dir
+    ):
+        paths = ["abc.123", "234.567", "d/object_1.pkl"]
+
+        if output_dir is not None:
+            args = paths + ["--output-dir", output_dir]
+        else:
+            args = paths
+
+        skops.io._cli.main_convert(command_line_args=args)
+        self.assert_called_correctly(mock_convert, paths=paths, output_dir=expected_dir)

skops/utils/entrypoint.py

@@ -0,0 +1,19 @@
+import argparse
+
+import skops.io._cli
+
+
+def main_entrypoint(command_line_args=None):
+    """Main entrypoint for all command line Skops methods."""
+    parser = argparse.ArgumentParser(
+        prog="Skops",
+        description="Main entrypoint for all command line Skops methods.",
+        add_help=False,
+    )
+
+    function_map = {"convert": skops.io._cli.main_convert}
+    parser.add_argument(
+        "function", help="Function to call.", choices=function_map.keys()
+    )
+    args, _ = parser.parse_known_args(command_line_args)
+    function_map.get(args.function)(parent=parser, command_line_args=command_line_args)

skops/utils/tests/test_entrypoint.py

@@ -0,0 +1,32 @@
+import pathlib
+import sys
+from unittest import mock
+
+import pytest
+
+from skops.utils.entrypoint import main_entrypoint
+
+
+class TestEntrypoint:
+    """Integration tests that check that entrypoint calls pass through correctly.
+
+    Full coverage of individual entrypoint calls should be done in their own classes.
+    """
+
+    @pytest.fixture(autouse=True)
+    def clear_argv(self):
+        # Required to clear argv in case Pytest is called on this specific function.
+        # Otherwise, clogs parser.parse_known_args() in argparse
+        sys.argv = [""]
+
+    @mock.patch("skops.io._cli._convert")
+    def test_convert_works_as_expected(self, mocked_convert: mock.MagicMock):
+        args = ["convert", "abc.def", "-t"]
+
+        main_entrypoint(args)
+
+        mocked_convert.assert_called_once_with(
+            input_file="abc.def",
+            output_dir=pathlib.Path.cwd(),
+            is_trusted=True,
+        )