Update write.py
This commit is contained in:
parent
4e7fc5b9b3
commit
6531bf7022
@ -1,9 +1,11 @@
|
||||
from typing import Union, Set, Optional, Dict
|
||||
from typing import Union, Set, Optional, Dict, DefaultDict
|
||||
from collections import defaultdict
|
||||
import traceback
|
||||
from peewee import (
|
||||
SqliteDatabase,
|
||||
MySQLDatabase,
|
||||
PostgresqlDatabase,
|
||||
Model
|
||||
)
|
||||
|
||||
from . import objects
|
||||
@ -41,6 +43,8 @@ class Session:
|
||||
self.added_album_ids: Dict[str] = dict()
|
||||
self.added_artist_ids: Dict[str] = dict()
|
||||
|
||||
self.db_objects: DefaultDict[data_models.BaseModel, list] = defaultdict(list)
|
||||
|
||||
def __enter__(self, database: Database):
|
||||
"""
|
||||
Enter the context of the database session
|
||||
@ -70,7 +74,7 @@ class Session:
|
||||
traceback.print_tb(exc_tb)
|
||||
print(f"Exception of type {exc_type} occurred with message: {exc_val}")
|
||||
|
||||
self.commit()
|
||||
self.commit(reset=False)
|
||||
return exc_val is None
|
||||
|
||||
def add_source(self, source: objects.Source, connected_to: data_models.Source.ContentTypes) -> data_models.Source:
|
||||
@ -81,6 +85,8 @@ class Session:
|
||||
content_object=connected_to
|
||||
).use(self.database)
|
||||
|
||||
self.db_objects[data_models.Source].append(db_source)
|
||||
|
||||
return db_source
|
||||
|
||||
def add_lyrics(self, lyrics: objects.Lyrics, song: data_models.Song) -> data_models.Lyrics:
|
||||
@ -91,6 +97,8 @@ class Session:
|
||||
song=song
|
||||
).use(self.database)
|
||||
|
||||
self.db_objects[data_models.Lyrics].append(db_lyrics)
|
||||
|
||||
for source in lyrics.source_list:
|
||||
self.add_source(source=source, connected_to=db_lyrics)
|
||||
|
||||
@ -104,6 +112,8 @@ class Session:
|
||||
song=song
|
||||
).use(self.database)
|
||||
|
||||
self.db_objects[data_models.Target].append(db_target)
|
||||
|
||||
return db_target
|
||||
|
||||
def add_song(self, song: objects.Song) -> Optional[data_models.Song]:
|
||||
@ -127,6 +137,7 @@ class Session:
|
||||
genre=song.genre
|
||||
).use(self.database)
|
||||
|
||||
self.db_objects[data_models.Song].append(db_song)
|
||||
self.added_song_ids[song.id] = db_song
|
||||
|
||||
for source in song.source_list:
|
||||
@ -171,6 +182,7 @@ class Session:
|
||||
|
||||
db_album = data_models.Album().use(self.database)
|
||||
|
||||
self.db_objects[data_models.Album].append(db_album)
|
||||
self.added_album_ids.add(album.id)
|
||||
|
||||
return db_album
|
||||
@ -189,12 +201,23 @@ class Session:
|
||||
|
||||
db_artist = data_models.Artist()
|
||||
|
||||
self.db_objects[data_models.Artist].append(db_artist)
|
||||
self.added_artist_ids[artist.id] = db_artist
|
||||
|
||||
return db_artist
|
||||
|
||||
def commit(self):
|
||||
def commit(self, reset: bool = True):
|
||||
"""
|
||||
Commit changes to the database
|
||||
"""
|
||||
pass
|
||||
|
||||
for model, model_instance_list in self.db_objects.items():
|
||||
model.Use(self.database).insert_many(model_instance_list)
|
||||
|
||||
if reset:
|
||||
self.__init__(self.database)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
with Session(SqliteDatabase(":memory:")) as session:
|
||||
session.add_song(objects.Song(title="Hs"))
|
Loading…
Reference in New Issue
Block a user