import glob import os import pooch class DownloadResource: """A remote resource that needs on demand downloading We use this as a wrapper to the pooch library. The wrapper registers each data file and allows prefetching through the CLI entry point misinformation_prefetch_models. """ # We store a list of defined resouces in a class variable, allowing # us prefetching from a CLI e.g. to bundle into a Docker image resources = [] def __init__(self, **kwargs): DownloadResource.resources.append(self) self.kwargs = kwargs def get(self): return pooch.retrieve(**self.kwargs) def misinformation_prefetch_models(): """Prefetch all the download resources""" for res in DownloadResource.resources: res.get() def find_files(path=None, pattern="*.png", recursive=True, limit=20): """Find image files on the file system :param path: The base directory where we are looking for the images. Defaults to None, which uses the XDG data directory if set or the current working directory otherwise. :param pattern: The naming pattern that the filename should match. Defaults to "*.png". Can be used to allow other patterns or to only include specific prefixes or suffixes. :param recursive: Whether to recurse into subdirectories. :param limit: The maximum number of images to be found. Defaults to 20. To return all images, set to None. """ if path is None: path = os.environ.get("XDG_DATA_HOME", ".") result = list(glob.glob(f"{path}/{pattern}", recursive=recursive)) if limit is not None: result = result[:limit] return result