Constrain python function arguments

Question

Is there a way in python to constrain an argument of a function to be only a set of values, the user should type the name of the column to change:

import sqlite3
def update_table(self, id: int, column_name: str, value: str):
    conn = sqlite3.connect("myDatabase.db")
    c = conn.cursor()
    c.execute("UPDATE table SET " + column_name + " = " + value +
              " WHERE id=" + id + ";")
    c.commit()
    c.close() 

id = input("Type the id: ")
c_name = input("Type the column name: ")
value = input("Type the new value: ")
update_table(id, c_name, value)

The problem is if the user type the name of a column that doesn't exist the code will break. What is the best way to constrain this choice?

It doesn't require a bunch of if's, you could just do `if(value in my_set)` or `assert value in my_set` — Spencer Wieczorek, Jun 02 '18 at 03:32
This is actually a better solution than what I was trying to do before, thanks. — Kostynha, Jun 02 '18 at 04:01

score 0 · Answer 1 · answered Jun 02 '18 at 10:52

Here is one option:

import sqlite3
def update_table(self, id: int, column_name: str, value: str):
    conn = sqlite3.connect("myDatabase.db")
    c = conn.cursor()
    c.execute("PRAGMA table_info(table_name);")
    col_names = set([col[1] for col in c.fetchall()])
    assert column_name in col_names, "{} is not a valid column name".format(column_name)
    c.execute("UPDATE table SET " + column_name + " = " + value +
              " WHERE id=" + id + ";")
    c.commit()
    c.close()

PRAGMA table_info(table_name); will return information about the table that looks something like this if you run a fetchall on it:

[(0, 'id', 'integer', 0, None, 1),
 (1, 'one', 'varchar(10)', 0, None, 0),
 (2, 'two', 'integer', 0, None, 0)]

Each tuple describes a column where index 1 of each column is the column name. So we create a set of possible column names to test against...

Constrain python function arguments

1 Answers1