Commit 1c57a059 authored by Stephan Seitz's avatar Stephan Seitz
Browse files

Assert same length when performing vector assignment

parent 38da1c39
Pipeline #21151 passed with stage
in 2 minutes and 25 seconds
......@@ -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)
import pytest
import sympy as sp
from pystencils import Assignment, AssignmentCollection
......@@ -52,6 +53,18 @@ def test_vector_assignments():
def test_wrong_vector_assignments():
"""From #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 ("""
......@@ -64,4 +77,3 @@ def test_vector_assignment_collection():
assignments = ps.AssignmentCollection([ps.Assignment(y,x)])
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment