import entities
from factories import BasicFactory

TAG_TO_ENTITY_MAPPING = {
    '100': entities.Person,
    '110': entities.Institution,
    '151': entities.Location
}


class INLFactory(BasicFactory):
    def __init__(self, tag_to_entity_mapping=None):
        self.mapping = tag_to_entity_mapping or TAG_TO_ENTITY_MAPPING

    def get_entity(self, entity_key, raw_object):
        if entity_key == '100':
            return entities.Person('', '', '')
        elif entity_key == '110':
            return entities.Institution()
        elif entity_key == '151':
            return entities.Location('', '', '')
        else:
            raise KeyError('Key {} was not recognized for factory {}'.format(entity_key, type(self)))