Made merge expressions function infer class wihthout warning

Author: RobbieKiwi

linopy/expressions.py

@@ -1941,50 +1941,43 @@ def merge(
     dim : str
         Dimension along which the expressions should be concatenated.
     cls : type
-        Type of the resulting expression.
+        Explicitly set the type of the resulting expression (So that the type checker will know the return type)
     **kwargs
         Additional keyword arguments passed to xarray.concat. Defaults to
         {coords: "minimal", compat: "override"} or, in the special case described
         above, to {coords: "minimal", compat: "override", "join": "override"}.
 
     Returns
     -------
-    res : linopy.LinearExpression
+    res : linopy.LinearExpression or linopy.QuadraticExpression
     """
-    if cls is None:
-        warn(
-            "Using merge without specifying the class is deprecated",
-            DeprecationWarning,
-        )
-        cls = LinearExpression
-
-    linopy_types = (variables.Variable, LinearExpression, QuadraticExpression)
-
     if not isinstance(exprs, list) and len(add_exprs):
         warn(
             "Passing a tuple to the merge function is deprecated. Please pass a list of objects to be merged",
             DeprecationWarning,
         )
         exprs = [exprs] + list(add_exprs)  # type: ignore
-    model = exprs[0].model
 
-    if (
-        cls is QuadraticExpression
-        and dim == TERM_DIM
-        and any(type(e) is LinearExpression for e in exprs)
-    ):
+    has_quad_expression = any(type(e) is QuadraticExpression for e in exprs)
+    has_linear_expression = any(type(e) is LinearExpression for e in exprs)
+    if cls is None:
+        cls = QuadraticExpression if has_quad_expression else LinearExpression
+
+    if cls is QuadraticExpression and dim == TERM_DIM and has_linear_expression:
         raise ValueError(
             "Cannot merge linear and quadratic expressions along term dimension."
             "Convert to QuadraticExpression first."
         )
 
-    if cls is not QuadraticExpression and any(
-        type(e) is QuadraticExpression for e in exprs
-    ):
+    if has_quad_expression and cls is not QuadraticExpression:
         raise ValueError(
-            "Cannot merge linear and quadratic expressions to QuadraticExpression"
+            "Cannot merge linear expressions to QuadraticExpression"
         )
 
+    linopy_types = (variables.Variable, LinearExpression, QuadraticExpression)
+
+    model = exprs[0].model
+
     if cls in linopy_types and dim in HELPER_DIMS:
         coord_dims = [
             {k: v for k, v in e.sizes.items() if k not in HELPER_DIMS} for e in exprs

test/test_linear_expression.py

@@ -1128,9 +1128,8 @@ def test_merge(x: Variable, y: Variable, z: Variable) -> None:
     res = merge([expr1, expr2], cls=LinearExpression)
     assert res.nterm == 6
 
-    with pytest.warns(DeprecationWarning):
-        res: LinearExpression = merge([expr1, expr2])  # type: ignore
-        assert res.nterm == 6
+    res: LinearExpression = merge([expr1, expr2])  # type: ignore
+    assert isinstance(res, LinearExpression)
 
     # now concat with same length of terms
     expr1 = z.sel(dim_0=0).sum("dim_1")

test/test_quadratic_expression.py

@@ -238,6 +238,9 @@ def test_merge_linear_expression_and_quadratic_expression(
     with pytest.raises(ValueError):
         merge([linexpr, quadexpr], cls=QuadraticExpression)
 
+    new_quad_ex = merge([linexpr.to_quadexpr(), quadexpr])  # type: ignore
+    assert isinstance(new_quad_ex, QuadraticExpression)
+
     with pytest.warns(DeprecationWarning):
         merge(quadexpr, quadexpr, cls=QuadraticExpression)  # type: ignore