From 46bd08ca87e610143672ebe392b8c3624afbe2d5 Mon Sep 17 00:00:00 2001 From: Yannik Schmidt Date: Sun, 25 Feb 2024 14:29:25 +0100 Subject: [PATCH] wip: ftp & progress --- data_backend.py | 89 ++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 88 insertions(+), 1 deletion(-) diff --git a/data_backend.py b/data_backend.py index e9c4f3d..6c881d6 100644 --- a/data_backend.py +++ b/data_backend.py @@ -2,17 +2,20 @@ 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, remote_root_dir=None): + 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): @@ -79,3 +82,87 @@ class LocalFS(DataBackend): 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