nwb-linkml/nwb_linkml/map.py

from dataclasses import dataclass
from typing import ClassVar, List, Optional
from enum import StrEnum
import ast
import re

class MAP_TYPES(StrEnum):
    key = 'key'
    """Mapping the name of one key to another key"""

class SCOPE_TYPES(StrEnum):
    namespace = 'namespace'

class PHASES(StrEnum):
    postload = "postload"
    """After the YAML for a model has been loaded"""


@dataclass
class Map:
    scope: str
    """The namespace that the map is relevant to"""
    scope_type: SCOPE_TYPES

    source: str
    """The path within the schema to select the element to transform"""
    target: str
    """The path where the element should end"""

    transform: Optional[callable] = None
    """
    Some transformation function, currently not implemented.
    """

    phase: Optional[PHASES] = None


    instances: ClassVar[List['Map']] = []
    """
    Maps that get defined!!!
    """

    def apply(self):
        raise NotImplementedError('do this in a subclass')

    def __post_init__(self):
        self.instances.append(self)



# def replace_keys(input: dict, source: str, target: str) -> dict:
#     """Recursively change keys in a dictionary"""


class KeyMap(Map):
    def apply(self, input: dict) -> dict:
        """
        Change all keys from source to target in a super naive way.

        Convert the dictionary to a string. Do regex. parse ast
        """
        input_str = str(input)
        input_str = re.sub(self.source, self.target, input_str)
        out = ast.literal_eval(input_str)
        return out


def apply_preload(ns_dict) -> dict:
    from nwb_linkml.maps import preload
    maps = [m for m in Map.instances if m.phase == PHASES.postload]
    for amap in maps:
        ns_dict = amap.apply(ns_dict)
    return ns_dict