FEAT CLI conversion

setup.py

@@ -34,6 +34,14 @@
 
 
 def setup_package():
+    package_data = dict(
+        entry_points={
+            "console_scripts": [
+                "skops = skops.cli.entrypoint:main_entrypoint",
+            ],
+        }
+    )
+
     metadata = dict(
         name=DISTNAME,
         maintainer=MAINTAINER,
@@ -72,7 +80,7 @@ def setup_package():
         include_package_data=True,
     )
 
-    setup(**metadata)
+    setup(**package_data, **metadata)
 
 
 if __name__ == "__main__":

skops/cli/__init__.py

@@ -0,0 +1,4 @@
+from ._convert import _convert_file, main_convert
+from .entrypoint import main_entrypoint
+
+__all__ = ["_convert_file", "main_convert", "main_entrypoint"]

skops/cli/_convert.py

@@ -0,0 +1,114 @@
+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_file(input_file: os.PathLike, output_file: pathlib.Path):
+    """
+    Function that is called by ``skops convert`` entrypoint.
+
+    Loads a pickle model from the input path, converts to skops format, and saves to
+    output file.
+
+    Parameters
+    ----------
+    input_file : os.PathLike
+        Path of input .pkl model to load.
+
+    output_file : os.PathLike
+        Path to save .skops model to.
+
+    """
+    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(
+            "Unknown Types Detected!\n"
+            f"While converting {model_name}, "
+            "the following unsafe types were found: \n"
+            f"{untrusted_str}\n\n"
+            f"When loading{output_file}, add ``--trusted`` to the skops.load call."
+        )
+
+    with open(output_file, "wb") as out_file:
+        logging.info(f"Writing to {output_file}")
+        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(
+        "-o",
+        "--output-files",
+        nargs="+",
+        help=(
+            "Specify output file names for the converted skops files."
+            "If not provided, will default to using the same name as the input file, "
+            "and saving to the current working directory."
+        ),
+        default=None,
+    )
+    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,
+    )
+
+    args = parser.parse_args(command_line_args)
+    output_files = args.output_files
+    input_files = args.inputs
+    logging.basicConfig(format="%(message)s", level=args.loglevel)
+
+    for input_file_index in range(len(args.inputs)):
+        input_file = input_files[input_file_index]
+        if output_files and len(output_files) > input_file_index:
+            output_file = output_files[input_file_index]
+        else:
+            # No filename provided, defaulting to base file path
+            file_name = pathlib.Path(input_file).stem
+            output_file = pathlib.Path.cwd() / f"{file_name}.skops"
+
+        _convert_file(
+            input_file=input_file,
+            output_file=output_file,
+        )

skops/cli/entrypoint.py

@@ -0,0 +1,19 @@
+import argparse
+
+import skops.cli._convert
+
+
+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.cli._convert.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/cli/tests/test_cli.py

@@ -0,0 +1,142 @@
+import logging
+import pathlib
+import pickle as pkl
+from typing import Optional
+from unittest import mock
+
+import numpy as np
+import pytest
+
+from skops import cli
+from skops.io import load
+
+
+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
+    ):
+        cli._convert_file(pkl_path, skops_path)
+        persisted_obj = 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)
+        cli._convert_file(pkl_path, skops_path)
+        persisted_obj = load(skops_path, trusted=True)
+
+        assert isinstance(persisted_obj, MockUnsafeType)
+
+        # 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_files: Optional[list] = None,
+    ):
+        if not output_files:
+            output_files = [
+                pathlib.Path.cwd() / f"{pathlib.Path(p).stem}.skops" for p in paths
+            ]
+        assert mock_convert.call_count == len(paths)
+        mock_convert.assert_has_calls(
+            [
+                mock.call(input_file=paths[i], output_file=output_files[i])
+                for i in range(len(paths))
+            ]
+        )
+
+    @mock.patch("skops.cli._convert._convert_file")
+    def test_base_works_as_expected(self, mock_convert: mock.MagicMock):
+        args = [
+            "123.pkl",
+            "abc.pkl",
+        ]
+
+        cli.main_convert(command_line_args=args)
+        self.assert_called_correctly(mock_convert, args)
+
+    @mock.patch("skops.cli._convert._convert_file")
+    @pytest.mark.parametrize(
+        "input_paths, output_files, expected_paths",
+        [
+            (["abc.123"], ["a/b/c"], ["a/b/c"]),
+            (["abc.123"], None, [pathlib.Path.cwd() / "abc.skops"]),
+        ],
+        ids=["Given an output path", "No output path"],
+    )
+    def test_with_output_dir_works_as_expected(
+        self, mock_convert: mock.MagicMock, input_paths, output_files, expected_paths
+    ):
+        if output_files is not None:
+            args = input_paths + ["--output-files"] + output_files
+        else:
+            args = input_paths
+
+        cli.main_convert(command_line_args=args)
+        self.assert_called_correctly(
+            mock_convert, paths=input_paths, output_files=expected_paths
+        )
+
+    @mock.patch("skops.cli._convert._convert_file")
+    @pytest.mark.parametrize(
+        "input_paths, output_files, expected_paths",
+        [
+            (
+                ["model_a.pkl", "model_b.pkl"],
+                ["a.skops", "b.skops"],
+                ["a.skops", "b.skops"],
+            ),
+            (
+                ["model_a.pkl", "model_b.pkl", "model_c.pkl"],
+                ["a.skops", "b.skops"],
+                ["a.skops", "b.skops", pathlib.Path.cwd() / "model_c.skops"],
+            ),
+        ],
+        ids=["With enough output paths", "With only some output paths"],
+    )
+    def test_for_multiple_inputs_and_outputs_works_as_expected(
+        self, mock_convert: mock.MagicMock, input_paths, output_files, expected_paths
+    ):
+        args = input_paths + ["--output-files"] + output_files
+        cli.main_convert(args)
+
+        self.assert_called_correctly(mock_convert, input_paths, expected_paths)

skops/cli/tests/test_entrypoint.py

@@ -0,0 +1,31 @@
+import pathlib
+import sys
+from unittest import mock
+
+import pytest
+
+from skops.cli.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.cli._convert._convert_file")
+    def test_convert_works_as_expected(self, mocked_convert: mock.MagicMock):
+        args = ["convert", "abc.def"]
+
+        main_entrypoint(args)
+
+        mocked_convert.assert_called_once_with(
+            input_file="abc.def",
+            output_file=pathlib.Path.cwd() / "abc.skops",
+        )