a bit of refactoring and implemented some more scraping

src/music_kraken/database/data_models.py

@@ -32,7 +32,7 @@ class BaseModel(Model):
     def Use(cls, database: Union[SqliteDatabase, PostgresqlDatabase, MySQLDatabase]) -> Model:
         cls._meta.database = database
         return cls
-    
+
     def use(self, database: Union[SqliteDatabase, PostgresqlDatabase, MySQLDatabase]) -> Model:
         self._meta.database = database
         return self
@@ -67,37 +67,6 @@ class Song(BaseModel):
     length: int = IntegerField(null=True)
     tracksort: int = IntegerField(null=True)
     genre: str = CharField(null=True)
-    # album: ForeignKeyField = ForeignKeyField(Album, backref='songs')
-
-
-class Source(BaseModel):
-    """A class representing a source of a song in the music database."""
-    ContentTypes = Union[Song, Album, Artist]
-
-    page: str = CharField()
-    url: str = CharField()
-
-    content_type: str = CharField()
-    content_id: int = IntegerField()
-    content: ForeignKeyField = ForeignKeyField('self', backref='content_items', null=True)
-
-    @property
-    def content_object(self) -> Union[Song, Album, Artist]:
-        """Get the content associated with the source as an object."""
-        if self.content_type == 'Song':
-            return Song.get(Song.id == self.content_id)
-        elif self.content_type == 'Album':
-            return Album.get(Album.id == self.content_id)
-        elif self.content_type == 'Artist':
-            return Artist.get(Artist.id == self.content_id)
-        else:
-            return None
-
-    @content_object.setter
-    def content_object(self, value: Union[Song, Album, Artist]) -> None:
-        """Set the content associated with the source as an object."""
-        self.content_type = value.__class__.__name__
-        self.content_id = value.id
 
 
 class Target(BaseModel):
@@ -131,9 +100,45 @@ class AlbumArtist(BaseModel):
     artist: ForeignKeyField = ForeignKeyField(Artist, backref='album_artists')
 
 
+class AlbumSong(BaseModel):
+    """A class representing the relationship between an album and an song."""
+    album: ForeignKeyField = ForeignKeyField(Album, backref='album_artists')
+    song: ForeignKeyField = ForeignKeyField(Song, backref='album_artists')
+
+
+class Source(BaseModel):
+    """A class representing a source of a song in the music database."""
+    ContentTypes = Union[Song, Album, Artist, Lyrics]
+
+    page: str = CharField()
+    url: str = CharField()
+
+    content_type: str = CharField()
+    content_id: int = IntegerField()
+    content: ForeignKeyField = ForeignKeyField('self', backref='content_items', null=True)
+
+    @property
+    def content_object(self) -> Union[Song, Album, Artist]:
+        """Get the content associated with the source as an object."""
+        if self.content_type == 'Song':
+            return Song.get(Song.id == self.content_id)
+        elif self.content_type == 'Album':
+            return Album.get(Album.id == self.content_id)
+        elif self.content_type == 'Artist':
+            return Artist.get(Artist.id == self.content_id)
+        else:
+            return None
+
+    @content_object.setter
+    def content_object(self, value: Union[Song, Album, Artist]) -> None:
+        """Set the content associated with the source as an object."""
+        self.content_type = value.__class__.__name__
+        self.content_id = value.id
+
+
 ALL_MODELS = [
-    Song, 
+    Song,
-    Album, 
+    Album,
     Artist,
     Source,
     Lyrics,
@@ -142,7 +147,6 @@ ALL_MODELS = [
     SongArtist
 ]
 
-
 if __name__ == "__main__":
     database_1 = SqliteDatabase(":memory:")
     database_1.create_tables([Song.Use(database_1)])

src/music_kraken/database/write.py

@@ -1,4 +1,4 @@
-from typing import Union, Set
+from typing import Union, Set, Optional, Dict
 import traceback
 from peewee import (
     SqliteDatabase,
@@ -12,6 +12,7 @@ from . import data_models
 # just a Type for type hintung. You can't do anything with it.
 Database = Union[SqliteDatabase, PostgresqlDatabase, MySQLDatabase]
 
+
 class Session:
     """
     Context manager for a database session
@@ -26,6 +27,7 @@ class Session:
     Attributes:
     database: An instance of a database connection from Peewee
     """
+
     def __init__(self, database: Database) -> None:
         """
         Initialize a database session
@@ -34,10 +36,10 @@ class Session:
         database: An instance of a database connection from Peewee
         """
         self.database = database
-        
+
-        self.added_song_ids: Set[str] = set()
+        self.added_song_ids: Dict[str] = dict()
-        self.added_album_ids: Set[str] = set()
+        self.added_album_ids: Dict[str] = dict()
-        self.added_artist_ids: Set[str] = set()
+        self.added_artist_ids: Dict[str] = dict()
 
     def __enter__(self, database: Database):
         """
@@ -49,7 +51,6 @@ class Session:
         Returns:
         self: The instance of the session object
         """
-        print('Entering the context')
         self.__init__(database=database)
         return self
 
@@ -71,9 +72,10 @@ class Session:
 
         self.commit()
         return exc_val is None
-    
+
     def add_source(self, source: objects.Source, connected_to: data_models.Source.ContentTypes) -> data_models.Source:
         db_source = data_models.Source(
+            id=source.id,
             page=source.page_str,
             url=source.url,
             content_object=connected_to
@@ -81,7 +83,30 @@ class Session:
 
         return db_source
 
-    def add_song(self, song: objects.Song) -> data_models.Song:
+    def add_lyrics(self, lyrics: objects.Lyrics, song: data_models.Song) -> data_models.Lyrics:
+        db_lyrics = data_models.Lyrics(
+            id=lyrics.id,
+            text=lyrics.text,
+            language=lyrics.language,
+            song=song
+        ).use(self.database)
+
+        for source in lyrics.source_list:
+            self.add_source(source=source, connected_to=db_lyrics)
+
+        return db_lyrics
+
+    def add_target(self, target: objects.Target, song: data_models.Song) -> data_models.Target:
+        db_target = data_models.Target(
+            id=target.id,
+            path=target.path,
+            file=target.file,
+            song=song
+        ).use(self.database)
+
+        return db_target
+
+    def add_song(self, song: objects.Song) -> Optional[data_models.Song]:
         """
         Add a song object to the session
 
@@ -91,8 +116,7 @@ class Session:
         if song.dynamic:
             return
         if song.id in self.added_song_ids:
-            return
+            return self.added_song_ids[song.id]
-        self.added_song_ids.add(song.id)
 
         db_song: data_models.Song = data_models.Song(
             id=song.id,
@@ -103,12 +127,37 @@ class Session:
             genre=song.genre
         ).use(self.database)
 
+        self.added_song_ids[song.id] = db_song
+
         for source in song.source_list:
             self.add_source(source=source, connected_to=db_song)
 
+        for target in [song.target]:
+            self.add_target(target, db_song)
+
+        for main_artist in song.main_artist_collection:
+            db_song_artist = data_models.SongArtist(
+                song=db_song,
+                artist=self.add_artist(main_artist),
+                is_feature=False
+            )
+
+        for feature_artist in song.feature_artist_collection:
+            db_song_artist = data_models.SongArtist(
+                song=db_song,
+                artist=self.add_artist(feature_artist),
+                is_feature=True
+            )
+
+        for album in [song.album]:
+            db_album_song = data_models.AlbumSong(
+                song=db_song,
+                album=self.add_album(album)
+            )
+
         return db_song
 
-    def add_album(self, album: objects.Album):
+    def add_album(self, album: objects.Album) -> Optional[data_models.Album]:
         """
         Add an album object to the session
 
@@ -118,10 +167,15 @@ class Session:
         if album.dynamic:
             return
         if album.id in self.added_album_ids:
-            return
+            return self.added_album_ids[album.id]
+
+        db_album = data_models.Album().use(self.database)
+
         self.added_album_ids.add(album.id)
 
-    def add_artist(self, artist: objects.Artist):
+        return db_album
+
+    def add_artist(self, artist: objects.Artist) -> Optional[data_models.Artist]:
         """
         Add an artist object to the session
 
@@ -131,11 +185,16 @@ class Session:
         if artist.dynamic:
             return
         if artist.id in self.added_artist_ids:
-            return
+            return self.added_artist_ids[artist.id]
-        self.added_artist_ids.add(artist.id)
+
+        db_artist = data_models.Artist()
+
+        self.added_artist_ids[artist.id] = db_artist
+
+        return db_artist
 
     def commit(self):
         """
         Commit changes to the database
         """
-        pass
+        pass