How to use 'assert' command to limit an input to only take number of columns in a set of d-CodePudding

I have written a function that can retrieve specific columns of data when given an input of the data and a variable called col_num but I need to write an assert command that does not let the input col_num exceed the number of columns in the data (For this I am assuming that I do not know the number of columns or there are too many to count).

My function is as follows:

import numpy as np
data = np.loadtxt('data.csv', delimiter = ',', skiprows = 1)
def get_column(data, col_num):
    assert(len(data.shape) == 2
    assert type(col_num) is int
    #assert goes here
    return data[:, col_num]

Any help would be greatly appreciated, thanks!

CodePudding user response：

Firstly, don't use assert for runtime checking. It can be optimised away: raise an exception instead.

Secondly, as the comments point out, you do have the data:

if data.shape[1] < col_num:
    raise ValueError(f"Supplied col num of {col_num} is greater than {data.shape[1]} columns")

CodePudding user response：

Not sure exactly what you're trying to achieve here. The assert will crash the program if the statement is False, but so would an uncaught IndexError. A benefit of relying on exceptions is that you can catch them further up the stack and gracefully handle runtime errors.

If you don't intend to do or log anything if an error occurs, just let the exception happen. If you need to log it or do something, catch the IndexError. If you can't fully handle it in the function, catch it, do what you can, then re-raise the exception:

import numpy as np
data = np.loadtxt('data.csv', delimiter = ',', skiprows = 1)
def get_column(data, col_num):
    try:
        return data[:, col_num]
    except IndexError as e:
        logging.exception(e)
        raise e