def _load_token_database(db, domain: Pattern[str]) -> tokens.Database:
"""Loads a Database from a database object, ELF, CSV, or binary database."""
if db is None:
return tokens.Database()
if isinstance(db, tokens.Database):
return db
if isinstance(db, elf_reader.Elf):
return _database_from_elf(db, domain)
# If it's a str, it might be a path. Check if it's an ELF or CSV.
if isinstance(db, (str, Path)):
if not os.path.exists(db):
raise FileNotFoundError(
f'"{db}" is not a path to a token database')
# Read the path as an ELF file.
with open(db, 'rb') as fd:
if elf_reader.compatible_file(fd):
return _database_from_elf(fd, domain)
# Read the path as a packed binary or CSV file.
return tokens.DatabaseFile(db)
# Assume that it's a file object and check if it's an ELF.
if elf_reader.compatible_file(db):
return _database_from_elf(db, domain)
# Read the database as CSV or packed binary from a file object's path.
if hasattr(db, 'name') and os.path.exists(db.name):
return tokens.DatabaseFile(db.name)
# Read CSV directly from the file object.
return tokens.Database(tokens.parse_csv(db))
def test_update_csv_file(self):
self._path.write_text(CSV_DATABASE)
db = tokens.DatabaseFile(self._path)
self.assertEqual(str(db), CSV_DATABASE)
db.add([tokens.TokenizedStringEntry(0xffffffff, 'New entry!')])
db.write_to_file()
self.assertEqual(self._path.read_text(),
CSV_DATABASE + 'ffffffff, ,"New entry!"\n')
def test_csv_not_utf8(self):
self._path.write_bytes(b'\x80' * 20)
with self.assertRaises(tokens.DatabaseFormatError):
tokens.DatabaseFile(self._path)
def test_csv_invalid_format_raises_exception(self):
self._path.write_text('MK34567890')
with self.assertRaises(tokens.DatabaseFormatError):
tokens.DatabaseFile(self._path)
def test_csv_file_too_short_raises_exception(self):
self._path.write_text('1234')
with self.assertRaises(tokens.DatabaseFormatError):
tokens.DatabaseFile(self._path)
