from typing import Literal

class WrapWorkers():
    def __init__(self):
        pass
    @staticmethod
    def wrap(worker_ids,data,target):
        input_args=(data,)
        output_args=(worker_ids,target,)
        return input_args,output_args

class WrapWorkersDiscriminator():
    def __init__(self):
        pass
    @staticmethod
    def wrap(worker_ids,data,target):
        input_args=(data,target)
        output_args=(worker_ids,)
        return input_args,output_args
    
def create_wrapper(wrap_method:Literal["default","discrim"]):
    if wrap_method=="default":
        wrapper=WrapWorkers()
    elif wrap_method=="discrim":
        wrapper=WrapWorkersDiscriminator()
    else:
        print("invalid wrap_method")
    return wrapper


class WrapNeglectId():
    def __init__(self):
        pass
    @staticmethod
    def wrap(worker_ids,data,target):
        return data,target
