Aluode/PerceptionLabPortable

1import asyncio2import functools3import inspect4 5import fsspec6from fsspec.asyn import AsyncFileSystem, running_async7 8from .chained import ChainedFileSystem9 10 11def async_wrapper(func, obj=None, semaphore=None):12    """13    Wraps a synchronous function to make it awaitable.14 15    Parameters16    ----------17    func : callable18        The synchronous function to wrap.19    obj : object, optional20        The instance to bind the function to, if applicable.21    semaphore : asyncio.Semaphore, optional22        A semaphore to limit concurrent calls.23 24    Returns25    -------26    coroutine27        An awaitable version of the function.28    """29 30    @functools.wraps(func)31    async def wrapper(*args, **kwargs):32        if semaphore:33            async with semaphore:34                return await asyncio.to_thread(func, *args, **kwargs)35        return await asyncio.to_thread(func, *args, **kwargs)36 37    return wrapper38 39 40class AsyncFileSystemWrapper(AsyncFileSystem, ChainedFileSystem):41    """42    A wrapper class to convert a synchronous filesystem into an asynchronous one.43 44    This class takes an existing synchronous filesystem implementation and wraps all45    its methods to provide an asynchronous interface.46 47    Parameters48    ----------49    sync_fs : AbstractFileSystem50        The synchronous filesystem instance to wrap.51    """52 53    protocol = "asyncwrapper", "async_wrapper"54    cachable = False55 56    def __init__(57        self,58        fs=None,59        asynchronous=None,60        target_protocol=None,61        target_options=None,62        semaphore=None,63        max_concurrent_tasks=None,64        **kwargs,65    ):66        if asynchronous is None:67            asynchronous = running_async()68        super().__init__(asynchronous=asynchronous, **kwargs)69        if fs is not None:70            self.sync_fs = fs71        else:72            self.sync_fs = fsspec.filesystem(target_protocol, **target_options)73        self.protocol = self.sync_fs.protocol74        self.semaphore = semaphore75        self._wrap_all_sync_methods()76 77    @property78    def fsid(self):79        return f"async_{self.sync_fs.fsid}"80 81    def _wrap_all_sync_methods(self):82        """83        Wrap all synchronous methods of the underlying filesystem with asynchronous versions.84        """85        excluded_methods = {"open"}86        for method_name in dir(self.sync_fs):87            if method_name.startswith("_") or method_name in excluded_methods:88                continue89 90            attr = inspect.getattr_static(self.sync_fs, method_name)91            if isinstance(attr, property):92                continue93 94            method = getattr(self.sync_fs, method_name)95            if callable(method) and not inspect.iscoroutinefunction(method):96                async_method = async_wrapper(method, obj=self, semaphore=self.semaphore)97                setattr(self, f"_{method_name}", async_method)98 99    @classmethod100    def wrap_class(cls, sync_fs_class):101        """102        Create a new class that can be used to instantiate an AsyncFileSystemWrapper103        with lazy instantiation of the underlying synchronous filesystem.104 105        Parameters106        ----------107        sync_fs_class : type108            The class of the synchronous filesystem to wrap.109 110        Returns111        -------112        type113            A new class that wraps the provided synchronous filesystem class.114        """115 116        class GeneratedAsyncFileSystemWrapper(cls):117            def __init__(self, *args, **kwargs):118                sync_fs = sync_fs_class(*args, **kwargs)119                super().__init__(sync_fs)120 121        GeneratedAsyncFileSystemWrapper.__name__ = (122            f"Async{sync_fs_class.__name__}Wrapper"123        )124        return GeneratedAsyncFileSystemWrapper125