Moving expression writer to util/misc

idaes/core/util/misc.py

@@ -15,9 +15,12 @@
 This module contains miscellaneous utility functions for use in IDAES models.
 """
 from enum import Enum
+import sys
 
 import pyomo.environ as pyo
 from pyomo.common.config import ConfigBlock
+from pyomo.core.expr.visitor import _ToStringVisitor
+from pyomo.core.base.expression import ExpressionData
 
 import idaes.logger as idaeslog
 
@@ -177,3 +180,66 @@ class StrEnum(str, Enum):
 
     def __str__(self):
         return str(self.value)
+
+
+class _ToExprStringVisitor(_ToStringVisitor):
+    """
+    Derived version fo the Pyomo _ToStringVisitor class
+    which checks for named Expressions in the expression tree
+    and prints their name instead of expanding the expression tree.
+    """
+
+    def visiting_potential_leaf(self, node):
+        # Check if node is a named Expression
+        if isinstance(node, ExpressionData):
+            # If it is, return the name of the Expression
+            return True, node.name
+
+        # Otherwise, continue descending as normal
+        return super().visiting_potential_leaf(node)
+
+
+def compact_expression_to_string(expr):
+    """Return a compact string representation of an expression.
+
+    Unlike the normal Pyomo string representations, this function will
+    identify any Expressions that appear within the expression tree and
+    print the name of these rather than expanding the expression tree.
+
+    Args:
+        expr: ExpressionBase. The root node of an expression tree.
+
+    Returns:
+        A string representation for the expression.
+    """
+    # Create and execute the visitor pattern
+    #
+    visitor = _ToExprStringVisitor(verbose=False, smap=None)
+    return visitor.dfs_postorder_stack(expr)
+
+
+def print_compact_form(expr, stream=None):
+    """
+    Writes a compact string representation fo the given component or
+    expression to the specified stream.
+
+    Unlike the normal Pyomo string representations, this function will
+    identify any Expressions that appear within the expression tree and
+    print the name of these rather than expanding the expression tree.
+
+    Args:
+        expr: component or expression to print
+        stream: StringIO object to write to. Default is stdout.
+
+    Returns:
+        None
+    """
+    if stream is None:
+        stream = sys.stdout
+
+    if hasattr(expr, "expr"):
+        # We have a Constraint or Expression
+        # We want to print the expression, not the object name
+        expr = expr.expr
+
+    stream.write(compact_expression_to_string(expr))

idaes/core/util/model_diagnostics.py

@@ -57,11 +57,9 @@
     NPV_UnaryFunctionExpression,
     NumericExpression,
 )
-from pyomo.core.expr.visitor import _ToStringVisitor
 from pyomo.core.base.block import BlockData
 from pyomo.core.base.var import VarData
 from pyomo.core.base.constraint import ConstraintData
-from pyomo.core.base.expression import ExpressionData
 from pyomo.repn.standard_repn import (  # pylint: disable=no-name-in-module
     generate_standard_repn,
 )
@@ -89,6 +87,7 @@
 from pyomo.core.base.units_container import _PyomoUnit
 
 from idaes.core.solvers.get_solver import get_solver
+from idaes.core.util.misc import compact_expression_to_string
 from idaes.core.util.model_statistics import (
     activated_blocks_set,
     deactivated_blocks_set,
@@ -4834,29 +4833,3 @@ def walk_expression(self, expr):
             const,
             self._cancellation_tripped,
         )
-
-
-class _ToExprStringVisitor(_ToStringVisitor):
-    def visiting_potential_leaf(self, node):
-        if isinstance(node, ExpressionData):
-            return True, node.name
-
-        return super().visiting_potential_leaf(node)
-
-
-def compact_expression_to_string(expr):
-    """Return a string representation of an expression.
-
-    Parameters
-    ----------
-    expr: ExpressionBase
-        The root node of an expression tree.
-
-    Returns:
-        A string representation for the expression.
-
-    """
-    # Create and execute the visitor pattern
-    #
-    visitor = _ToExprStringVisitor(verbose=False, smap=None)
-    return visitor.dfs_postorder_stack(expr)

idaes/core/util/tests/test_misc.py

@@ -14,16 +14,19 @@
 This module contains miscellaneous utility functions for use in IDAES models.
 """
 
+from io import StringIO
 import pytest
 import re
 
-from pyomo.environ import ConcreteModel, Set, Block, Var, units
+from pyomo.environ import ConcreteModel, Constraint, Expression, Set, Block, Var, units
 from pyomo.common.config import ConfigBlock
 from pyomo.core.base.units_container import UnitsError
 
 from idaes.core.util.misc import (
     add_object_reference,
     set_param_from_config,
+    compact_expression_to_string,
+    print_compact_form,
 )
 import idaes.logger as idaeslog
 
@@ -361,3 +364,54 @@ def test_unitted_default(self, caplog):
             "b no units provided for parameter test_param - assuming "
             "default units" in caplog.text
         )
+
+
+class TestToExprStringVisitor:
+    @pytest.mark.unit
+    def test_compact_expression_to_string(self):
+        m = ConcreteModel()
+        m.v1 = Var()
+        m.v2 = Var()
+        m.v3 = Var()
+
+        m.e1 = Expression(expr=m.v1 + m.v2)
+        m.c1 = Constraint(expr=0 == m.v3 * m.e1)
+
+        str = compact_expression_to_string(m.c1.expr)
+        expected = "v3*e1  ==  0"
+
+        assert str == expected
+
+    @pytest.mark.unit
+    def test_print_compact_form_expr(self):
+        m = ConcreteModel()
+        m.v1 = Var()
+        m.v2 = Var()
+        m.v3 = Var()
+
+        m.e1 = Expression(expr=m.v1 + m.v2)
+        m.c1 = Constraint(expr=0 == m.v3 * m.e1)
+
+        expected = "v3*e1  ==  0"
+
+        stream = StringIO()
+        print_compact_form(m.c1.expr, stream=stream)
+
+        assert stream.getvalue() == expected
+
+    @pytest.mark.unit
+    def test_print_compact_form_component(self):
+        m = ConcreteModel()
+        m.v1 = Var()
+        m.v2 = Var()
+        m.v3 = Var()
+
+        m.e1 = Expression(expr=m.v1 + m.v2)
+        m.c1 = Constraint(expr=0 == m.v3 * m.e1)
+
+        expected = "v3*e1  ==  0"
+
+        stream = StringIO()
+        print_compact_form(m.c1, stream=stream)
+
+        assert stream.getvalue() == expected

idaes/core/util/tests/test_model_diagnostics.py

@@ -97,7 +97,6 @@
     check_parallel_jacobian,
     compute_ill_conditioning_certificate,
     ConstraintTermAnalysisVisitor,
-    compact_expression_to_string,
 )
 from idaes.core.util.parameter_sweep import (
     SequentialSweepRunner,
@@ -5375,20 +5374,3 @@ def test_mole_fraction_constraint_trace(self):
         assert len(cc) == 0
         assert not k
         assert not tr
-
-
-class TestToExprStringVisitor:
-    @pytest.mark.unit
-    def test_visitor(self):
-        m = ConcreteModel()
-        m.v1 = Var()
-        m.v2 = Var()
-        m.v3 = Var()
-
-        m.e1 = Expression(expr=m.v1 + m.v2)
-        m.c1 = Constraint(expr=0 == m.v3 * m.e1)
-
-        str = compact_expression_to_string(m.c1.expr)
-        expected = "v3*e1  ==  0"
-
-        assert str == expected