/
database.py
156 lines (121 loc) · 3.75 KB
/
database.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
'''
database.py
This file contains a simple JSON storage class.
'''
import json
import os
import classifier
import soundfiles
import fingerprint
import colors
c = colors.Colors()
class Database:
'''A simple JSON storage system.'''
DBPREFIX = 'databases/'
def __init__(self, name, replace=False, read=True):
'''
Database initializer. Reads database if exists.
Args:
name: name of database.
replace: whether to replace or append database.
read: whether to load the databse from disk on init.
'''
self.replace = replace
self.read = read
self.name = self.DBPREFIX + name
if read:
self.database = self._read_db()
else:
self.database = []
def add(self, tokens):
'''
Pushes a singleton or list of tokens to the database.
Args:
tokens: Token or list of Token.
'''
if not isinstance(tokens, list):
tokens = [tokens]
for token in tokens:
self._push_token(token)
def save(self):
'''Wrapper for internal write.'''
self._write_db()
def load(self):
'''Wrapper for internal read.'''
self.database = self._read_db()
def as_classifier(self):
'''
Build a Classifier from database.
Returns:
Classifier populated with Tokens from database.
'''
cl = classifier.Classifier()
for entry in self.database:
cl.add_token(
classifier.Token(
entry['fingerprint'],
entry['time'],
entry['filename']
)
)
return cl
def populate(self, directory):
'''
Analyzes all audio files in a directory and add their fingerprint
to the database. Writes result to disk when done.
Args:
directory: directory containing audio files.
'''
audio_files = sorted(soundfiles.find_files(directory + '/*'))
c.notice("Reading {} files...".format(len(audio_files)))
i = 0
for filename in audio_files:
i += 1
c.write('\r\t-file ' + str(i) + ' of ' + str(len(audio_files)))
signal = soundfiles.load_signal(filename)
self.add(fingerprint.get_tokens(signal))
c.write("\n")
c.notice("Writing database to disk...")
self.save()
c.succes("Database written")
def get_size(self):
'''Returns current size of database'''
return len(self.database)
def _exists(self):
'''Returns whether the current database exists.'''
return os.path.isfile(self.name)
def _remove(self):
'''Removes the database file if exists.'''
if self._exists():
os.remove(self.name)
def _push_token(self, token):
'''
Pushes a token to the database.
Args:
token: Token.
'''
self.database.append(token.as_dict())
def _read_db(self):
'''
Read the stored database JSON.
Returns:
JSON content of file.
'''
if not self._exists():
return []
with open(self.name, "r") as file:
try:
return json.load(file)
except ValueError:
return []
def _write_db(self):
'''Write database to disk.'''
if not os.path.exists(self.DBPREFIX):
os.makedirs(self.DBPREFIX)
if self._exists() and self.replace:
self._remove()
with open(self.name, 'w+') as file:
json.dump(self.database, file)
def __str__(self):
'''Returns database name.'''
return self.name