Querying named tuples

# -*- coding: utf-8 -*- from operator import lt, le, eq, ne, ge, gt from collections import namedtuple, Counter fields = 'name gender age city state smoking' Person = namedtuple('Person', fields) fields_arr = fields.split(' ') func_dict = {'<': lt, '<=': le, '=': eq, '!=': ne, '>=': ge, '>': gt} def query(selects, from_iterable, wheres): select_idxs, where_idxs = [], [] conds, where_vals = [], [] wheres_is_dict = isinstance(wheres, dict) if isinstance(wheres, str): wheres = [wheres] if isinstance(selects, str): selects = [selects] if wheres_is_dict: # Construct string queries with equal sign wheres_list = [] for k, v in wheres.items(): for vsub in v: wheres_list.append('%s = %s' % (k, vsub)) wheres = wheres_list wheres_len = len(wheres) for select in selects: if select in fields_arr: select_idxs.append(fields_arr.index(select)) for where in wheres: attr, cond, where_val = where.split(' ', 2); if attr in fields_arr: where_idxs.append(fields_arr.index(attr)) conds.append(cond) where_vals.append(where_val) if len(select_idxs) == 0 or len(where_idxs) == 0: return 'Unable to tell from the data.' res = [] for item in from_iterable: conditions_met = 0 for i, where_idx in enumerate(where_idxs): cond, where_val = conds[i], where_vals[i] f = func_dict[cond] val = float(where_val) if where_val.isdigit() else where_val if f(item[where_idx], val): conditions_met += 1 if wheres_len == conditions_met or (wheres_is_dict and conditions_met): sub_res = [] for select_idx in select_idxs: sub_res.append(str(item[select_idx])) res.append(' '.join(sub_res)) return res people = [ Person('John', 'male', 35, 'Ottawa', 'Canada', 'smoking'), Person('Sophia', 'female', 27, 'Melbourne', 'Australia', 'not smoking'), Person('Enrique', 'male', 24, 'Madrid', 'Spain', 'not smoking'), Person('Lena', 'female', 31, 'Stuttgart', 'Germany', 'not smoking'), Person('Florian', 'male', 26, 'Hamburg', 'Germany', 'smoking'), Person('Denise', 'female', 28, 'Paris', 'France', 'smoking'), Person('Julieta', 'female', 27, 'La Plata', 'Argentina', 'smoking'), Person('Gabriela', 'female', 26, 'Vienna', 'Austria', 'not smoking'), Person('Ingmar', 'male', 29, 'Malmö', 'Sweden', 'not smoking'), Person('Juha', 'male', 26, 'Turku', 'Finnland', 'not smoking'), Person('Hugo', 'male', 28, 'Paris', 'France', 'smoking'), Person('Benoit', 'male', 29, 'Montpellier', 'France', 'not smoking'), Person('Adrian', 'male', 28, 'Bucharest', 'Romania', 'smoking'), Person('Jaromir', 'male', 28, 'Wrocław', 'Poland', 'smoking'), Person('Thiago', 'male', 32, 'São Paulo', 'Brazil', 'not smoking'), Person('Charlotte', 'female', 28, 'Manchester', 'Great Britain', 'not smoking') ] # Get the names of all people older than 28 print(query('name', people, 'age > 28')) # ['John', 'Lena', 'Ingmar', 'Benoit', 'Thiago'] # Get the names of all people whose age is between 24 and 26 print(query('name', people, ['age >= 24', 'age <= 26'])) # ['Enrique', 'Florian', 'Gabriela', 'Juha'] # Select the cities of all people named Jaromir print(query('city', people, 'name = Jaromir')) # ['Wrocław'] # Select the names of all people from Germany print(query('name', people, 'state = Germany')) # ['Lena', 'Florian'] # Select the age, gender and name of all people from La Plata, Argentina print(query(['age', 'gender', 'name'], people, ['city = La Plata', 'state = Argentina'])) # ['27 female Julieta'] # How many cigarette smokers are male/female? print(dict(Counter(query('gender', people, 'smoking = smoking')))) # {'male': 5, 'female': 2} # Does Gabriela smoke? print(query('smoking', people, 'name = Gabriela')) # ['not smoking'] # Where does Juha come from? print(query('state', people, 'name = Juha')) # ['Finnland'] # Do people from Malmö use bicycles? print(query('cycling', people, 'city = Malmö')) # Unable to tell from the data. # List the names of all people from France, Germany, Poland or Romania in reverse sorted order q = query('name', people, {'state': ['France', 'Germany', 'Poland', 'Romania']}) print(sorted(q, reverse=True)) # ['Lena', 'Jaromir', 'Hugo', 'Florian', 'Denise', 'Benoit', 'Adrian'] # or q.sort() print(q[::-1]) # ['Lena', 'Jaromir', 'Hugo', 'Florian', 'Denise', 'Benoit', 'Adrian']