Add array/set comprehension support

Implement array and set comprehensions in Expression: add helpers to validate and render generator clauses (_generator_clause_to_mz), convert comprehension expressions (_comprehension_expr_to_mz), and public factories array_comprehension and set_comprehension. These methods validate inputs, accept callables that are invoked with generator variables, support optional generator predicates (rendered with "where"), and raise existing Pymzm exceptions on invalid input. Tests updated to cover rendering, use in model constraints, and input validation for comprehensions.

Author: Arthod

src/pymzm/expression.py

@@ -219,6 +219,69 @@ def _predicate_to_mz(var_name: str, predicate) -> str:
 
         raise PymzmValueIsNotCondition("predicate", predicate)
 
+    @staticmethod
+    def _generator_clause_to_mz(generator, arg_name: str="generators") -> str:
+        if (not isinstance(generator, (tuple, list)) or len(generator) not in (2, 3)):
+            raise PymzmValueIsNotExpression(arg_name, generator)
+
+        var_name = generator[0]
+        domain = generator[1]
+        predicate = generator[2] if (len(generator) == 3) else None
+
+        if (not isinstance(var_name, str) or not var_name.strip()):
+            raise PymzmValueIsNotExpression("var_name", var_name)
+
+        domain_mz = Expression._domain_to_mz(domain)
+        clause = f"{var_name} in {domain_mz}"
+        if (predicate is not None):
+            predicate_mz = Expression._predicate_to_mz(var_name, predicate)
+            clause += f" where {predicate_mz}"
+
+        return clause
+
+    @staticmethod
+    def _comprehension_expr_to_mz(expr, generator_var_names):
+        if (isinstance(expr, Callable)):
+            vars_expr = [Expression(name) for name in generator_var_names]
+            if (len(vars_expr) == 1):
+                expr = expr(vars_expr[0])
+            else:
+                expr = expr(*vars_expr)
+
+        if (isinstance(expr, bool)):
+            return "true" if expr else "false"
+
+        if (isinstance(expr, (Expression, int, float))):
+            return str(expr)
+
+        raise PymzmValueIsNotExpression("expr", expr)
+
+    @staticmethod
+    def array_comprehension(expr, generators):
+        if (not isinstance(generators, Iterable) or isinstance(generators, (str, bytes))):
+            raise PymzmValueIsNotExpression("generators", generators)
+
+        generators = list(generators)
+        if (not len(generators)):
+            raise PymzmNoValues("generators")
+
+        clauses = [Expression._generator_clause_to_mz(g) for g in generators]
+        expr_mz = Expression._comprehension_expr_to_mz(expr, [g[0] for g in generators])
+        return Expression(f"[{expr_mz} | {', '.join(clauses)}]")
+
+    @staticmethod
+    def set_comprehension(expr, generators):
+        if (not isinstance(generators, Iterable) or isinstance(generators, (str, bytes))):
+            raise PymzmValueIsNotExpression("generators", generators)
+
+        generators = list(generators)
+        if (not len(generators)):
+            raise PymzmNoValues("generators")
+
+        clauses = [Expression._generator_clause_to_mz(g) for g in generators]
+        expr_mz = Expression._comprehension_expr_to_mz(expr, [g[0] for g in generators])
+        return Expression(f"{{{expr_mz} | {', '.join(clauses)}}}")
+
     @staticmethod
     def predicate(name: str, *args) -> "ExpressionBool":
         if (not isinstance(name, str) or not name.strip()):

tests/test_expression_validation.py

@@ -128,6 +128,45 @@ def test_expression_render_is_valid_mzn_when_embedded_in_model(self):
 
         assert_valid_mzn(self, model.model_mzn_str)
 
+    def test_array_and_set_comprehension_render_with_filters(self):
+        arr = pymzm.Expression.array_comprehension(
+            lambda i: i + 1,
+            [("i", range(1, 4), lambda i: i >= 2)],
+        )
+        s = pymzm.Expression.set_comprehension(
+            lambda i, j: i + j,
+            [("i", "1..2"), ("j", [1, 2, 3], lambda j: j >= 2)],
+        )
+
+        self.assertEqual(str(arr), "[(i + 1) | i in 1..3 where (i >= 2)]")
+        self.assertEqual(str(s), "{(i + j) | i in 1..2, j in {1, 2, 3} where (j >= 2)}")
+
+    def test_comprehensions_can_be_used_in_model_constraints(self):
+        model = pymzm.Model()
+        x = model.add_variable("x", val_min=0, val_max=10)
+
+        arr = pymzm.Expression.array_comprehension(lambda i: i + x, [("i", range(1, 3))])
+        s = pymzm.Expression.set_comprehension(lambda i: i, [("i", [1, 2, 3], lambda i: i >= 2)])
+
+        model.add_constraint(x >= 0)
+        model.add_constraint(pymzm.Expression._func("sum", [arr]) >= 0)
+        model.add_constraint(pymzm.Expression._func("card", [s]) >= 1)
+        model.set_solve_criteria(pymzm.SOLVE_SATISFY)
+        model.generate()
+
+        self.assertIn("sum([", model.model_mzn_str)
+        self.assertIn("{i | i in {1, 2, 3} where (i >= 2)}", model.model_mzn_str)
+        assert_valid_mzn(self, model.model_mzn_str)
+
+    def test_comprehensions_validate_inputs(self):
+        self.assertRaises(pymzm.PymzmNoValues, pymzm.Expression.array_comprehension, 1, [])
+        self.assertRaises(pymzm.PymzmValueIsNotExpression, pymzm.Expression.array_comprehension, object(), [("i", [1])])
+        self.assertRaises(pymzm.PymzmValueIsNotExpression, pymzm.Expression.array_comprehension, 1, ["bad"])
+        self.assertRaises(pymzm.PymzmValueIsNotCondition, pymzm.Expression.array_comprehension, 1, [("i", [1], 7)])
+
+        self.assertRaises(pymzm.PymzmNoValues, pymzm.Expression.set_comprehension, 1, [])
+        self.assertRaises(pymzm.PymzmValueIsNotExpression, pymzm.Expression.set_comprehension, object(), [("i", [1])])
+
 
 if __name__ == "__main__":
     unittest.main()