backend/section_infer_helper/base_helper.py

from backend.utils.data_process import LANGUAGE_EXT_MAP

from abc import ABC, abstractmethod

class BaseHelper(ABC):
    
    def _get_lang_ext(language_list):
        ext_list = []
        for lang in language_list:
            ext_list.extend(LANGUAGE_EXT_MAP.get(lang, []))
        return ext_list
    
    
    def _get_lang_by_ext(ext):
        for lang, ext_list in LANGUAGE_EXT_MAP.items():
            if ext in ext_list:
                return lang
        return None
    
    
    class InputData():
        def __init__(self, filename, patch, section, patch_msg):
            self.filename = filename
            self.patch = patch
            self.section = section
            self.patch_msg = patch_msg
    
    
    @abstractmethod
    def load_model(self, *args, **kwargs):
        raise NotImplementedError()
    
    
    @abstractmethod
    def infer(self, diff_code, message = None, batch_size = 1):
        '''
        Result format:
        [
            file: [
                {
                    "section": section,
                    "predict": 1/0,
                    "conf": conf
                },
                ...
            ],
            ...
        ]
        '''
        raise NotImplementedError()