file_experiment_repository.py

import json
import os
from typing import Any, Dict, Optional, Set

import pandas as pd

from sziszapangma.integration.repository.experiment_repository import ExperimentRepository


class FileExperimentRepository(ExperimentRepository):
    _metadata_file_path: str
    _pretty_format: bool
    _cache_value: Any

    def __init__(self, metadata_file_path: str, pretty_format: bool = False):
        super().__init__()
        self._metadata_file_path = metadata_file_path
        self._pretty_format = pretty_format
        self._cache_value = None

    def initialise(self):
        if not os.path.isfile(self._metadata_file_path):
            with open(self._metadata_file_path, "w") as f:
                empty_dict = dict()
                f.write(json.dumps(empty_dict))

    def property_exists(self, record_id: str, property_name: str) -> bool:
        json_content = self._get_file_parsed_json()
        if record_id not in json_content:
            return False
        elif property_name not in json_content[record_id]:
            return False
        else:
            return True

    def update_property_for_key(self, record_id: str, property_name: str, property_value: Any):
        json_content = self._get_file_parsed_json()
        if record_id not in json_content:
            json_content[record_id] = dict({})
        json_content[record_id][property_name] = property_value
        self._update_metadata_file(json_content)

    def get_property_for_key(self, record_id: str, property_name: str) -> Optional[Any]:
        json_content = self._get_file_parsed_json()
        if self.property_exists(record_id, property_name):
            return json_content[record_id][property_name]
        else:
            return None

    def _get_file_parsed_json(self) -> Dict[str, Dict[Any, Any]]:
        if self._cache_value is None:
            with open(self._metadata_file_path, "r") as f:
                self._cache_value = json.loads(f.read())
        return self._cache_value

    def _update_metadata_file(self, json_content: Dict[str, Any]):
        self._cache_value = json_content
        indent = 4 if self._pretty_format else None
        with open(self._metadata_file_path, "w") as f:
            f.write(json.dumps(json_content, indent=indent))

    def delete_property_for_key(self, record_id: str, property_name: str):
        json_content = self._get_file_parsed_json()
        if record_id not in json_content:
            json_content[record_id] = dict({})
        json_content[record_id].pop(property_name)
        self._update_metadata_file(json_content)

    def get_all_record_ids(self) -> Set[str]:
        return set(self._get_file_parsed_json().keys())

    def get_all_properties(self) -> Set[str]:
        properties = set()
        all_metadata = self._get_file_parsed_json()
        for key in all_metadata.keys():
            for prop in all_metadata[key].keys():
                properties.add(prop)
        return properties