Merge pull request #29 from spellew/additional-entity-support

Implemented retrieval and/or testing of additional entities

brainzutils/musicbrainz_db/artist.py

@@ -8,21 +8,19 @@
 from brainzutils.musicbrainz_db.includes import check_includes
 
 
-def get_artist_by_id(mbid):
+def get_artist_by_id(mbid, includes=None):
     """Get artist with MusicBrainz ID.
     Args:
         mbid (uuid): MBID(gid) of the artist.
     Returns:
         Dictionary containing the artist information
     """
-    artist = _get_artist_by_id(mbid)
-    return artist
-
+    if includes is None:
+        includes = []
 
-def _get_artist_by_id(mbid):
     return fetch_multiple_artists(
         [mbid],
-        includes=['artist-rels', 'url-rels'],
+        includes=includes,
     ).get(mbid)
 
 

brainzutils/musicbrainz_db/editor.py

@@ -7,21 +7,19 @@
 from brainzutils.musicbrainz_db.includes import check_includes
 
 
-def get_editor_by_id(editor_id):
+def get_editor_by_id(editor_id, includes=None):
     """Get editor with editor ID.
     Args:
         editor_id (int): ID of the editor.
     Returns:
         Dictionary containing the editor information
     """
-    editor = _get_editor_by_id(editor_id)
-    return editor
-
+    if includes is None:
+        includes = []
 
-def _get_editor_by_id(editor_id):
     return fetch_multiple_editors(
         [editor_id],
-        includes=[],
+        includes=includes,
     ).get(editor_id)
 
 
@@ -35,6 +33,7 @@ def fetch_multiple_editors(editor_ids, includes=None):
     """
     if includes is None:
         includes = []
+
     includes_data = defaultdict(dict)
     check_includes('editor', includes)
     with mb_session() as db:

brainzutils/musicbrainz_db/event.py

@@ -0,0 +1,90 @@
+from collections import defaultdict
+from mbdata import models
+from brainzutils.musicbrainz_db import mb_session
+from brainzutils.musicbrainz_db.utils import get_entities_by_gids
+from brainzutils.musicbrainz_db.includes import check_includes
+from brainzutils.musicbrainz_db.serialize import serialize_events
+from brainzutils.musicbrainz_db.helpers import get_relationship_info
+
+
+def get_event_by_id(mbid, includes=None):
+    """Get event with the MusicBrainz ID.
+
+    Args:
+        mbid (uuid): MBID(gid) of the event.
+    Returns:
+        Dictionary containing the event information.
+    """
+    if includes is None:
+        includes = []
+
+    return fetch_multiple_events(
+        [mbid],
+        includes=includes,
+    ).get(mbid)
+
+
+def fetch_multiple_events(mbids, includes=None):
+    """Get info related to multiple events using their MusicBrainz IDs.
+
+    Args:
+        mbids (list): List of MBIDs of events.
+        includes (list): List of information to be included.
+
+    Returns:
+        Dictionary containing info of multiple events keyed by their mbid.
+    """
+    if includes is None:
+        includes = []
+    includes_data = defaultdict(dict)
+    check_includes('event', includes)
+    with mb_session() as db:
+        query = db.query(models.Event)
+        events = get_entities_by_gids(
+            query=query,
+            entity_type='event',
+            mbids=mbids,
+        )
+        event_ids = [event.id for event in events.values()]
+
+        if 'artist-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='artist',
+                source_type='event',
+                source_entity_ids=event_ids,
+                includes_data=includes_data,
+            )
+        if 'place-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='place',
+                source_type='event',
+                source_entity_ids=event_ids,
+                includes_data=includes_data,
+            )
+        if 'series-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='series',
+                source_type='event',
+                source_entity_ids=event_ids,
+                includes_data=includes_data,
+            )
+        if 'url-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='url',
+                source_type='event',
+                source_entity_ids=event_ids,
+                includes_data=includes_data,
+            )
+        if 'release-group-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='release_group',
+                source_type='event',
+                source_entity_ids=event_ids,
+                includes_data=includes_data,
+            )
+    return {str(mbid): serialize_events(events[mbid], includes_data[events[mbid].id]) for mbid in mbids}

brainzutils/musicbrainz_db/helpers.py

@@ -59,20 +59,21 @@ def _relationship_link_helper(relation, query, source_attr, target_attr, target_
             setdefault(relation_type, []).append(link)
 
 
-def get_tags(db, entity_model, tag_model, entity_ids):
+def get_tags(db, entity_model, tag_model, foreign_tag_id, entity_ids):
     """Get tags associated with entities.
 
     Args:
         db (Session object): Session object.
         entity_model (mbdata.models): Model of the entity.
         tag_model (mbdata.models): Tag of the model.
+        foreign_tag_id (tag_model.foreign_key): Foreign ID that joins the tag model and entity model
         entity_ids (list): IDs of the entity whose tags are to be fetched
 
     Returns:
         List of tuples containing the entity_ids and the list of associated tags.
     """
     tags = db.query(entity_model.id, func.array_agg(Tag.name)).\
-        join(tag_model).\
+        join(tag_model, entity_model.id == foreign_tag_id).\
         join(Tag).\
         filter(entity_model.id.in_(entity_ids)).\
         group_by(entity_model.id).\

brainzutils/musicbrainz_db/includes.py

@@ -31,6 +31,8 @@
         "artists", "labels", "recordings", "release-groups", "media", "annotation", "aliases"
     ] + TAG_INCLUDES + RELATION_INCLUDES,
     'artist': ["recordings", "releases", "media", "aliases", "annotation"] + RELATION_INCLUDES + TAG_INCLUDES,
+    'label': ["area", "aliases", "annotation"] + RELATION_INCLUDES + TAG_INCLUDES,
+    'work': ["artists", "recordings", "aliases", "annotation"] + RELATION_INCLUDES + TAG_INCLUDES,
     'editor': [],  # TODO: List includes here (BU-18)
 }
 

brainzutils/musicbrainz_db/label.py

@@ -0,0 +1,74 @@
+from collections import defaultdict
+from mbdata import models
+from sqlalchemy.orm import joinedload
+from brainzutils.musicbrainz_db import mb_session
+from brainzutils.musicbrainz_db.utils import get_entities_by_gids
+from brainzutils.musicbrainz_db.includes import check_includes
+from brainzutils.musicbrainz_db.serialize import serialize_labels
+from brainzutils.musicbrainz_db.helpers import get_relationship_info
+
+
+def get_label_by_id(mbid, includes=None):
+    """Get label with the MusicBrainz ID.
+
+    Args:
+        mbid (uuid): MBID(gid) of the label.
+    Returns:
+        Dictionary containing the label information.
+    """
+    if includes is None:
+        includes = []
+
+    return fetch_multiple_labels(
+        [mbid],
+        includes=includes,
+    ).get(mbid)
+
+
+def fetch_multiple_labels(mbids, includes=None):
+    """Get info related to multiple labels using their MusicBrainz IDs.
+
+    Args:
+        mbids (list): List of MBIDs of labels.
+        includes (list): List of information to be included.
+    Returns:
+        Dictionary containing info of multiple labels keyed by their mbid.
+    """
+    if includes is None:
+        includes = []
+    includes_data = defaultdict(dict)
+    check_includes('label', includes)
+    with mb_session() as db:
+        query = db.query(models.Label).\
+            options(joinedload("type")).\
+            options(joinedload("area"))
+        labels = get_entities_by_gids(
+            query=query,
+            entity_type='label',
+            mbids=mbids,
+        )
+        label_ids = [label.id for label in labels.values()]
+
+        if 'artist-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='artist',
+                source_type='label',
+                source_entity_ids=label_ids,
+                includes_data=includes_data,
+            )
+
+        if 'url-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='url',
+                source_type='label',
+                source_entity_ids=label_ids,
+                includes_data=includes_data,
+            )
+
+        for label in labels.values():
+            includes_data[label.id]['type'] = label.type
+            includes_data[label.id]['area'] = label.area
+
+    return {str(mbid): serialize_labels(labels[mbid], includes_data[labels[mbid].id]) for mbid in mbids}

brainzutils/musicbrainz_db/place.py

@@ -0,0 +1,82 @@
+from collections import defaultdict
+from mbdata import models
+from sqlalchemy.orm import joinedload
+from brainzutils.musicbrainz_db import mb_session
+from brainzutils.musicbrainz_db.includes import check_includes
+from brainzutils.musicbrainz_db.serialize import serialize_places
+from brainzutils.musicbrainz_db.helpers import get_relationship_info
+from brainzutils.musicbrainz_db.utils import get_entities_by_gids
+
+
+def get_place_by_id(mbid, includes=None):
+    """Get place with the MusicBrainz ID.
+
+    Args:
+        mbid (uuid): MBID(gid) of the place.
+    Returns:
+        Dictionary containing the place information.
+    """
+    if includes is None:
+        includes = []
+
+    return fetch_multiple_places(
+        [mbid],
+        includes=includes,
+    ).get(mbid)
+
+
+def fetch_multiple_places(mbids, includes=None):
+    """Get info related to multiple places using their MusicBrainz IDs.
+
+    Args:
+        mbids (list): List of MBIDs of places.
+        includes (list): List of information to be included.
+
+    Returns:
+        Dictionary containing info of multiple places keyed by their mbid.
+    """
+    if includes is None:
+        includes = []
+    includes_data = defaultdict(dict)
+    check_includes('place', includes)
+    with mb_session() as db:
+        query = db.query(models.Place).\
+            options(joinedload("area")).\
+            options(joinedload("type"))
+        places = get_entities_by_gids(
+            query=query,
+            entity_type='place',
+            mbids=mbids,
+        )
+        place_ids = [place.id for place in places.values()]
+
+        if 'artist-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='artist',
+                source_type='place',
+                source_entity_ids=place_ids,
+                includes_data=includes_data,
+            )
+        if 'place-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='place',
+                source_type='place',
+                source_entity_ids=place_ids,
+                includes_data=includes_data,
+            )
+        if 'url-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='url',
+                source_type='place',
+                source_entity_ids=place_ids,
+                includes_data=includes_data,
+            )
+
+        for place in places.values():
+            includes_data[place.id]['area'] = place.area
+            includes_data[place.id]['type'] = place.type
+        places = {str(mbid): serialize_places(places[mbid], includes_data[places[mbid].id]) for mbid in mbids}
+    return places

brainzutils/musicbrainz_db/recording.py

@@ -20,8 +20,11 @@ def get_recording_by_mbid(mbid, includes=None):
     """
     if includes is None:
         includes = []
-    recording = get_many_recordings_by_mbid([mbid], includes)
-    return recording.get(mbid)
+    return fetch_multiple_recordings(
+        [mbid],
+        includes=includes,
+    ).get(mbid)
+
 
 def get_many_recordings_by_mbid(mbids, includes=None):
     """ Get multiple recordings with MusicBrainz IDs. It fetches recordings
@@ -37,12 +40,10 @@ def get_many_recordings_by_mbid(mbids, includes=None):
     if includes is None:
         includes = []
 
-    recordings = _fetch_multiple_recordings(mbids, includes)
-
-    return recordings
+    return fetch_multiple_recordings(mbids, includes)
 
 
-def _fetch_multiple_recordings(mbids, includes=None):
+def fetch_multiple_recordings(mbids, includes=None):
     """ Fetch multiple recordings with MusicBrainz IDs.
 
     Args:
@@ -89,6 +90,24 @@ def _fetch_multiple_recordings(mbids, includes=None):
             for recording in recordings.values():
                 includes_data[recording.id]['artists'] = recording.artist_credit.artists
 
+        if 'url-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='url',
+                source_type='recording',
+                source_entity_ids=recording_ids,
+                includes_data=includes_data,
+            )
+
+        if 'work-rels' in includes:
+            get_relationship_info(
+                db=db,
+                target_type='work',
+                source_type='recording',
+                source_entity_ids=recording_ids,
+                includes_data=includes_data,
+            )
+
         serial_recordings = {str(mbid): serialize_recording(recordings[mbid], includes_data[recordings[mbid].id]) for mbid in mbids}
 
     return serial_recordings