black & ruff updates

Signed-off-by: Łukasz Ćmielowski <Lukasz.Cmielowski@pl.ibm.com>
onnx · Dec 20, 2024 · 67ab5bf · 67ab5bf
1 parent 3f834fa
commit 67ab5bf
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/onnxmltools/convert/coreml/shape_calculators/OneHotEncoder.py b/onnxmltools/convert/coreml/shape_calculators/OneHotEncoder.py
@@ -26,7 +26,7 @@ def calculate_one_hot_encoder_output_shapes(operator):
         operator.outputs[0].type = FloatTensorType(
             [N, len(int_categories)], doc_string=operator.outputs[0].type.doc_string
         )
-    elif len(str_categories) > 0 and type(operator.inputs[0].type) == StringTensorType:
+    elif len(str_categories) > 0 and type(operator.inputs[0].type) is StringTensorType:
         operator.outputs[0].type = FloatTensorType(
             [N, len(str_categories)], doc_string=operator.outputs[0].type.doc_string
         )

diff --git a/onnxmltools/convert/lightgbm/_parse.py b/onnxmltools/convert/lightgbm/_parse.py
@@ -36,7 +36,9 @@ def __init__(self, booster):
         elif self.objective_.startswith("multiclass"):
             self.operator_name = "LgbmClassifier"
             self.classes_ = self._generate_classes(booster)
-        elif self.objective_.startswith(("regression", "poisson", "gamma", "quantile", "huber")):
+        elif self.objective_.startswith(
+            ("regression", "poisson", "gamma", "quantile", "huber")
+        ):
             self.operator_name = "LgbmRegressor"
         else:
             raise NotImplementedError(

diff --git a/tests/lightgbm/test_objective_functions.py b/tests/lightgbm/test_objective_functions.py
@@ -31,7 +31,13 @@
 
 
 class ObjectiveTest(unittest.TestCase):
-    _regressor_objectives: Tuple[str] = ("regression", "poisson", "gamma", "quantile", "huber")
+    _regressor_objectives: Tuple[str] = (
+        "regression",
+        "poisson",
+        "gamma",
+        "quantile",
+        "huber",
+    )
 
     @staticmethod
     def _calc_initial_types(X: DataFrame) -> List[Tuple[str, TensorType]]: