import numbers


def diag(v, k=0):
    """Create a diagonal matrix from a vector or return a vector of the diagonal of a matrix."""

    is_vector = isinstance(v[0], numbers.Number)
    size = len(v)
    d = []

    if is_vector:
        # Create a diagonal matrix with the provided values
        for i, value in enumerate(v):
            # Check that the matrix is square, we .cannot invert the matrix if it is not
            d.append([0] * i + [value] + [0] * (size - i - 1))
    else:  # pragma: no cover
        for r in v:
            if len(r) != size:
                raise ValueError('Matrix must be a n x n matrix')
            if 0 <= k < size:
                d.append(r[k])
            k += 1
    return d