set_magnitude gives a division by zero if v is a zero vector:
def set_magnitude(v, mag):
"""
return: Vector v with magnitude set to mag.
"""
scale = mag / magnitude(v)
return mult(v, scale)
Could either change the above to the following or make it a separate function set_magnitude_zero():
old_mag = magnitude(v)
if old_mag > 0.0:
return mult(v, mag / old_mag)
return copy.copy(v)