import os import glob import yaml import software import ftplib import tqdm class DataBackend: def _create_cache_dir(self, cache_dir): os.makedirs(cache_dir, exist_ok=True) def __init__(self, user, password, install_dir, server=None, remote_root_dir=None): self.user = user self.password = password self.remote_root_dir = remote_root_dir self.server = server self.install_dir = install_dir def get(self, path, return_content=False): '''Return the contents of this path''' raise NotImplementedError() def list(self, path): '''List the contents of this path''' raise NotImplementedError() def find_all_metadata(self): '''Return key-value map of { software : metadata-dict }''' raise NotImplementedError class LocalFS(DataBackend): def get(self, path, cache_dir=None, return_content=False): # check the load cache dir # if cache_dir: self._create_cache_dir(cache_dir) elif not cache_dir and not return_content: AssertionError("Need to set either cache_dir or return_content!") # prepend root dir if not given # fullpath = path if self.remote_root_dir and not path.startswith(self.remote_root_dir): fullpath = os.path.join(self.remote_root_dir, path) # load the file on remote # with open(fullpath, "rb") as f: print(cache_dir, path) target = os.path.join(cache_dir, os.path.basename(path)) with open(target, "wb") as ft: if return_content: return f.read() ft.write(f.read()) return target def list(self, path, fullpaths=False): # prepend root dir if not given # fullpath = path if self.remote_root_dir and not path.startswith(self.remote_root_dir): fullpath = os.path.join(self.remote_root_dir, path) if not os.path.isdir(fullpath): return [] if fullpaths: return [ os.path.join(path, filename) for filename in os.listdir(fullpath)] else: return os.listdir(fullpath) def find_all_metadata(self): meta_info_list = [] for software_dir in glob.iglob(self.remote_root_dir + "/*"): meta_file = os.path.join(software_dir, "meta.yaml") if not os.path.isfile(meta_file): continue else: meta_info_list.append(software.Software(meta_file, self)) return meta_info_list class FTP(DataBackend): def _connect(self): if self.server.startswith("ftp://"): ftp = ftplib.FTP(self.server) elif self.server.startswith("ftps://"): ftp = ftplib.FTP_TLS(self.server) if self.user: ftp.login(self.user, self.password) else: ftp.login() return ftp def get(self, path, cache_dir=None, return_content=False): # prepend root dir if not given # fullpath = path if self.remote_root_dir and not path.startswith(self.remote_root_dir): fullpath = os.path.join(self.remote_root_dir, path) ftp = self._connect() # load the file on remote # total_size = ftp.size(fullpath) local_file = os.path.join(cache_dir, os.path.basename(path)) self.progress_bar["maximum"] = total_size with open(local_file, 'wb') as local_file, tqdm( desc="Downloading", total=total_size, unit='B', unit_scale=True ) as progress_bar: # Define a callback function to update the progress bar # def callback(data): local_file.write(data) self.root.update_idletasks() # Update the GUI self.progress_bar.step(len(data)) # run with callback # ftp.retrbinary('RETR ' + fullpath, callback) with open(fullpath, "rb") as f: print(cache_dir, path) target = os.path.join(cache_dir, os.path.basename(path)) with open(target, "wb") as ft: if return_content: return f.read() ft.write(f.read()) return target def list(self, path, fullpaths=False): # prepend root dir if not given # fullpath = path if self.remote_root_dir and not path.startswith(self.remote_root_dir): fullpath = os.path.join(self.remote_root_dir, path) if not os.path.isdir(fullpath): return [] if fullpaths: return [ os.path.join(path, filename) for filename in os.listdir(fullpath)] else: return os.listdir(fullpath) def find_all_metadata(self): meta_info_list = [] for software_dir in glob.iglob(self.remote_root_dir + "/*"): meta_file = os.path.join(software_dir, "meta.yaml") if not os.path.isfile(meta_file): continue else: meta_info_list.append(software.Software(meta_file, self)) return meta_info_list