Assert same length when performing vector assignment

pystencils/assignment.py

@@ -26,7 +26,8 @@ if Assignment:
     _old_new = sp.codegen.ast.Assignment.__new__
 
     def _Assignment__new__(cls, lhs, rhs, *args, **kwargs):
-        if isinstance(lhs, (list, set, tuple, sp.Matrix)) and isinstance(rhs, (list, set, tuple, sp.Matrix)):
+        if isinstance(lhs, (list, tuple, sp.Matrix)) and isinstance(rhs, (list, tuple, sp.Matrix)):
+            assert len(lhs) == len(rhs), f'{lhs} and {rhs} must have same length when performing vector assignment!'
             return tuple(_old_new(cls, a, b, *args, **kwargs) for a, b in zip(lhs, rhs))
         return _old_new(cls, lhs, rhs, *args, **kwargs)
 

pystencils_tests/test_assignment_collection.py

+import pytest
 import sympy as sp
 
 from pystencils import Assignment, AssignmentCollection
@@ -52,6 +53,18 @@ def test_vector_assignments():
     print(assignments)
 
 
+def test_wrong_vector_assignments():
+    """From #17 (https://i10git.cs.fau.de/pycodegen/pystencils/issues/17)"""
+
+    import pystencils as ps
+    import sympy as sp
+    a, b = sp.symbols("a b")
+    
+    with pytest.raises(AssertionError,
+            match=r'Matrix(.*) and Matrix(.*) must have same length when performing vector assignment!'):
+        ps.Assignment(sp.Matrix([a,b]), sp.Matrix([1,2,3]))
+
+
 def test_vector_assignment_collection():
     """From #17 (https://i10git.cs.fau.de/pycodegen/pystencils/issues/17)"""
 
@@ -64,4 +77,3 @@ def test_vector_assignment_collection():
 
     assignments = ps.AssignmentCollection([ps.Assignment(y,x)])
     print(assignments)
-