open2c / cooltools

The tools for your .cool's
MIT License
138 stars 51 forks source link

Pool decorator doesn't allow 3rd party multiprocessing map functors #535

Closed dmitrymyl closed 1 week ago

dmitrymyl commented 2 months ago

Hey! I tried to run cooltools.pileup with a custom map_functor, but it threw an error: image

It seems the issue is with the cooltools.lib.common.pool_decorator, specifically on this line and the last line in the code snippet below:

    @wraps(func)
    def wrapper(*args, **kwargs):
        # If alternative or third party map functors are provided
        if "map_functor" in kwargs.keys():
            logging.info(f"using an alternative map functor: {kwargs['map_functor']}")
            return func(*args, **kwargs, map_functor=kwargs["map_functor"])

"map_functor" is passed twice to the func: once in **kwargs and next in map_functor=kwargs["map_functor"], and that's why it breaks.

To fix that, I suggest a patch to the last line, so that it would look like:

    @wraps(func)
    def wrapper(*args, **kwargs):
        # If alternative or third party map functors are provided
        if "map_functor" in kwargs.keys():
            logging.info(f"using an alternative map functor: {kwargs['map_functor']}")
            return func(*args, **kwargs)

As a workaround before the patch, I am using cooltools.pileup.__wrapped__, which strips the pool_decorator away.

Hope it will help improve the package!