MonomialTermExpression.creaate_node_with_local_data: resolve custom t…

…ype error
jsiirola · Oct 3, 2024 · 4a5e593 · 4a5e593
1 parent 91023a1
commit 4a5e593
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/pyomo/core/expr/numeric_expr.py b/pyomo/core/expr/numeric_expr.py
@@ -1094,7 +1094,7 @@ def create_node_with_local_data(self, args, classtype=None):
             # types, the simplest / fastest thing to do is just defer to
             # the operator dispatcher.
             return operator.mul(*args)
-        return self.__class__(args)
+        return classtype(args)
 
 
 class DivisionExpression(NumericExpression):

diff --git a/pyomo/core/tests/unit/test_numeric_expr.py b/pyomo/core/tests/unit/test_numeric_expr.py
@@ -4313,6 +4313,17 @@ def test_sin(self):
             total = counter.count - start
             self.assertEqual(total, 1)
 
+    def test_create_node_with_local_data(self):
+        e = self.m.p * self.m.a
+        self.assertIs(type(e), MonomialTermExpression)
+
+        f = e.create_node_with_local_data([self.m.b, self.m.p])
+        self.assertIs(type(f), MonomialTermExpression)
+        self.assertStructuredAlmostEqual(f._args_, [self.m.p, self.m.b])
+
+        g = e.create_node_with_local_data([self.m.b, self.m.p], ProductExpression)
+        self.assertIs(type(g), ProductExpression)
+        self.assertStructuredAlmostEqual(g._args_, [self.m.b, self.m.p])
 
 #
 # Fixed               - Expr has a fixed value