lint

slim_trees/__init__.py

@@ -42,7 +42,9 @@
 
 
 def dump_sklearn_compressed(
-    model: Any, file: Union[str, Path, BinaryIO], compression: Optional[Union[str, dict]] = None
+    model: Any,
+    file: Union[str, Path, BinaryIO],
+    compression: Optional[Union[str, dict]] = None,
 ):
     """
     Pickles a model and saves a compressed version to the disk.
@@ -79,7 +81,9 @@ def dumps_sklearn_compressed(
 
 
 def dump_lgbm_compressed(
-    model: Any, file: Union[str, Path, BinaryIO], compression: Optional[Union[str, dict]] = None
+    model: Any,
+    file: Union[str, Path, BinaryIO],
+    compression: Optional[Union[str, dict]] = None,
 ):
     """
     Pickles a model and saves a compressed version to the disk.

slim_trees/sklearn_tree.py

@@ -46,7 +46,7 @@ def dumps(model: Any) -> bytes:
 def _tree_pickle(tree: Tree):
     assert isinstance(tree, Tree)
     reconstructor, args, state = tree.__reduce__()
-    compressed_state = _compress_tree_state(state) # type: ignore
+    compressed_state = _compress_tree_state(state)  # type: ignore
     return _tree_unpickle, (reconstructor, args, (slim_trees_version, compressed_state))
 
 
@@ -113,7 +113,7 @@ def _compress_tree_state(state: Dict) -> Dict:
             "values": values,
         },
         **(
-            {"missing_go_to_left": np.packbits(missing_go_to_left)} # type: ignore
+            {"missing_go_to_left": np.packbits(missing_go_to_left)}  # type: ignore
             if sklearn_version_ge_130
             else {}
         ),

tests/test_lgbm_compression.py

@@ -115,14 +115,11 @@ def test_dump_and_load_from_file(tmp_path, lgbm_regressor):
         load_compressed(file, compression="lzma")
 
     # No compression method specified
-    with pytest.raises(ValueError), (tmp_path / "model.pickle.lzma").open(
-        "rb"
-    ) as file:
+    with pytest.raises(ValueError), (tmp_path / "model.pickle.lzma").open("rb") as file:
         load_compressed(file)
 
-    with pytest.raises(ValueError), (tmp_path / "model.pickle.lzma").open(
-        "wb"
-    ) as file:
+    with pytest.raises(ValueError), (tmp_path / "model.pickle.lzma").open("wb") as file:
         dump_lgbm_compressed(lgbm_regressor, file)
 
+
 # todo add tests for large models

tests/test_sklearn_compression.py

@@ -167,13 +167,11 @@ def test_dump_and_load_from_file(tmp_path, random_forest_regressor):
         load_compressed(file, compression="lzma")
 
     # No compression method specified
-    with pytest.raises(ValueError), (tmp_path / "model.pickle.lzma").open(
-        "rb"
-    ) as file:
+    with pytest.raises(ValueError), (tmp_path / "model.pickle.lzma").open("rb") as file:
         load_compressed(file)
 
-    with pytest.raises(ValueError), (tmp_path / "model.pickle.lzma").open(
-        "wb"
-    ) as file:
+    with pytest.raises(ValueError), (tmp_path / "model.pickle.lzma").open("wb") as file:
         dump_sklearn_compressed(random_forest_regressor, file)
+
+
 # todo add tests for large models