I'm looking for a neat way to return multiple matrices in python

Question

I'm trying to find a neater way of returning an arbitrary number of NxN matrices. Right now I'm using this function

from numpy import matrix, zeros

def empty_matrix( dim, num ):   
"""Returns an empty square matrix of type complex and size a."""
    if num == 1:
        return matrix(zeros(shape =( dim, dim ), dtype=complex))
    else:
        return [ matrix(zeros(shape =( dim, dim ), dtype=complex)) for _ in range( num )]

in the following way:

A,B,C = empty_matrix( 2, 3 ) # sets A, B, C as 3 2x2 matrices

I'm trying to figure out if there's a way to avoid the if/else statement. Any ideas?

Answer 1

Drop the if num == 1, and use a 1-element-tuple-unpacking:

A,B,C = empty_matrices( 2, 3 ) # sets A, B, C as 3 2x2 matrices
A, = empty_matrices( 2, 1 ) # sets A as a 2x2 matrix