diff --git a/soitool/database.py b/soitool/database.py index f7131aa8396512c4a587ff1cdd2fc2006ed13733..3aeeb76eafa899b7ebfa2db9a056b368a815e42f 100644 --- a/soitool/database.py +++ b/soitool/database.py @@ -120,25 +120,33 @@ class Database(): return categories - def get_small_codebook(self): + def get_codebook(self, small=False): """ - Retrive the entries belonging to the small codebook. + Retrive the entries belonging to the small codebook or big codebook. + + Parameters + ---------- + small : Bool + Full or smal codebook to be returned Returns ------- - Lsit of dicts + codebook : list (of dicts) [{'word': str, 'type': str, 'category': str, 'code': str}] - - """ - stmt = 'SELECT * FROM CodeBook WHERE Type = ?' - queried = self.conn.execute(stmt, (1,)).fetchall() + # Get either small or whole codebook + stmt = 'SELECT * FROM CodeBook' + if (small): + stmt = stmt + ' WHERE Type = ?' + queried = self.conn.execute(stmt, (1,)).fetchall() + else: + queried = self.conn.execute(stmt).fetchall() - small_codebook = [] + codebook = [] for entry in queried: - small_codebook.append({'word': entry['Word'], - 'category': entry['Category'], - 'type': entry['Type'], - 'code': entry['Code']}) + codebook.append({'word': entry['Word'], + 'category': entry['Category'], + 'type': entry['Type'], + 'code': entry['Code']}) - return small_codebook + return codebook