PyCS/models/fixed_model/Pipeline.py

from json import dump, load
from os import path
from typing import List

from pycs.pipeline.Fit import Fit
from pycs.pipeline.Job import Job
from pycs.pipeline.Pipeline import Pipeline as Interface


class Pipeline(Interface):
    def __init__(self, root_folder, distribution):
        print('fmv1 init')
        self.root_folder = root_folder

    def close(self):
        print('fmv1 close')

    def execute(self, job: Job) -> List[dict]:
        print('fmv1 execute')
        data_file = path.join(self.root_folder, 'data.json')

        if path.exists(data_file):
            with open(data_file, 'r') as file:
                result = load(file)
        else:
            result = {}

        if job.path in result:
            return result[job.path]
        else:
            return []

    def fit(self, fit: List[Fit]):
        print('fmv1 fit')
        result = {}

        for f in fit:
            result[f.path] = f.result

        data_file = path.join(self.root_folder, 'data.json')
        with open(data_file, 'w') as file:
            dump(result, file, indent=4)