Merge pull request #139 from c-hydro/dev

Dev

Author: ltrotter

pyproject.toml

@@ -38,8 +38,7 @@ dependencies = [
     "pyhdf>=0.10.2",           #only for modis_downloader
     "cdsapi>=0.7.4",           #only for cds_downloader
     "ecmwf-opendata >= 0.2.0", #only for ecmwf_opendata_downloader
-    "paramiko>=2.9.3",
-    "ftpretty>=0.1.0", # I would like to avoid a repetition here, paramiko and ftpretty kind of do the same thing
+    "paramiko>=2.9.3"
 ]
 
 [project.urls]

src/door/base_downloaders.py

@@ -7,7 +7,11 @@
 import xarray as xr
 import os
 
-from .utils.io import download_http, check_download, handle_missing, download_ftp, download_sftp
+import paramiko
+import ftplib
+import requests
+
+from .utils.io import check_download, handle_missing
 
 from d3tools import spatial as sp
 from d3tools import timestepping as ts
@@ -164,13 +168,13 @@ def _loop_timesteps_and_save_data(self, timesteps: list[ts.TimeStep]) -> None:
         # the latter is more space efficient, but at times you have to download the data for several timesteps at once
         # so it is better to have the option to download all the data in a single folder
         if self.single_temp_folder:
-            with tempfile.TemporaryDirectory(ignore_cleanup_errors=True, prefix=os.getenv('TMP')) as tmp_path:
+            with tempfile.TemporaryDirectory(ignore_cleanup_errors=True, dir=os.getenv('TMP')) as tmp_path:
                 for timestep in timesteps:
                     self._get_and_save_data_ts(timestep, tmp_path)
                     rm_at_exit(tmp_path)
         else:
             for timestep in timesteps:
-                with tempfile.TemporaryDirectory(ignore_cleanup_errors=True, prefix=os.getenv('TMP')) as tmp_path:
+                with tempfile.TemporaryDirectory(ignore_cleanup_errors=True, dir=os.getenv('TMP')) as tmp_path:
                     self._get_and_save_data_ts(timestep, tmp_path)
                     rm_at_exit(tmp_path)
 
@@ -349,20 +353,14 @@ class URLDownloader(DOORDownloader):
 
     name = "URL_Downloader"
 
-    def __init__(self, url_blank: str, protocol: str = 'http', host: str|None = None) -> None:
+    def __init__(self, url_blank: str, protocol: str = 'http') -> None:
 
         self.url_blank = url_blank
-        if protocol.lower() not in ['http', 'ftp', 'sftp', 'https']:
+        if protocol.lower() not in ['http', 'https']:
             raise ValueError(f'Protocol {protocol} not supported')
         else:
             self.protocol = protocol.lower()
 
-        if self.protocol == 'ftp' or self.protocol == 'sftp':
-            if host is None:
-                raise ValueError(f'FTP host must be specified')
-            else:
-                self.host = host
-
         super().__init__()
 
     def format_url(self, **kwargs) -> str:
@@ -382,14 +380,15 @@ def download(self, destination: str, min_size: float = None, missing_action: str
 
         url = self.format_url(**kwargs)
         try:
-            if self.protocol == 'http' or self.protocol == 'https':
-                download_http(url, destination, kwargs["auth"])
-            elif self.protocol == 'ftp':
-                download_ftp(self.host, url, destination, kwargs["auth"])
-            elif self.protocol == 'sftp':
-                download_sftp(self.host, url, destination, kwargs["auth"])
-            else:
-                raise ValueError(f'Protocol {self.protocol} not supported')
+            r = requests.get(url, kwargs["auth"])
+            if r.status_code != 200:
+                raise FileNotFoundError(r.text)
+            
+            os.makedirs(os.path.dirname(destination), exist_ok=True)
+
+            with open(destination, 'wb') as f:
+                f.write(r.content)
+
         except Exception as e:
             handle_missing(missing_action, kwargs)
             self.log.debug(f'Error downloading {url}: {e}')
@@ -403,6 +402,98 @@ def download(self, destination: str, min_size: float = None, missing_action: str
 
         return True
 
+class FTPDownloader(DOORDownloader):
+    """
+    Downloader for data from an FTP server via FTP or SFTP.
+    This typer of downloader is useful for data that can be downloaded from an FTP server.
+    It allows to specify a URL template with placeholders for various parameters (as keyword arguments).
+    """
+
+    name = "FTP_Downloader"
+
+    def __init__(self, host: str, port: int = 21, protocol: str = 'ftp', user: str = 'anonymous', password: str = 'anonymous') -> None:
+        if protocol.lower() not in ['ftp', 'sftp']:
+            raise ValueError(f'Protocol {protocol} not supported')
+        self.protocol = protocol.lower()
+        
+        self.host = host
+        self.port = port
+        self.user = user
+        self.password = password
+        super().__init__()
+
+        if self.protocol == 'sftp':
+            self.transport = paramiko.Transport((host, port))
+            self.transport.connect(username=user, password=password)
+            self.client = paramiko.SFTPClient.from_transport(self.transport)
+        elif self.protocol == 'ftp':
+            self.client = ftplib.FTP()
+            self.client.connect(host, port)
+            self.client.login(user, password)
+
+    def __del__(self):
+        """
+        Close the FTP or SFTP client connection when the downloader is deleted.
+        """
+        if hasattr(self, 'client'):
+            try:
+                self.client.close()
+            except Exception as e:
+                self.log.debug(f'Error closing {self.protocol} client: {e}')
+
+        if hasattr(self, 'transport'):
+            try:
+                self.transport.close()
+            except Exception as e:
+                self.log.debug(f'Error closing {self.protocol} transport: {e}')
+        
+    def download(self, blank_path, destination: str, min_size: float = None, missing_action: str = 'error', **kwargs) -> bool:
+        """
+        Downloads data from FTP or SFTP server.
+        Eventually check file size to avoid empty files.
+        """
+
+        url = blank_path.format(**kwargs)
+
+        try:
+            if self.protocol == 'sftp':
+                self.client.get(url, destination)
+            elif self.protocol == 'ftp':
+                with open(destination, 'wb') as f:
+                    self.client.retrbinary(f'RETR {url}', f.write)
+        except Exception as e:
+            handle_missing(missing_action, kwargs)
+            self.log.debug(f'Error downloading {url} via {self.protocol}: {e}')
+            return False
+
+        success_flag, success_msg = check_download(destination, min_size, missing_action)
+        if success_flag > 0:
+            handle_missing(missing_action, kwargs)
+            self.log.debug(f'Error downloading file from {url}: {success_msg}')
+            return False
+
+        return True
+
+    def check_data(self, blank_path,  **kwargs) -> bool:
+        """
+        Check if the data is available on the FTP or SFTP server.
+        This method can be used to check if the data is available before downloading it.
+        """
+        url = blank_path.format(**kwargs)
+        if self.protocol == 'ftp':
+            # For FTP, we can use the 'nlst' command to check if the file exists
+            if len(self.client.nlst(url)) > 0:
+                return True
+        elif self.protocol == 'sftp':
+            try:
+                self.client.stat(url)
+                return True
+            except Exception as e:
+                pass
+        return False
+                
+            
+
 class APIDownloader(DOORDownloader):
     """
     Downloader for data from an API.
@@ -443,4 +534,4 @@ def download(self, destination: str, min_size: float = None, missing_action: str
     def retrieve(self, **kwargs):
         return self.client.retrieve(**kwargs)
 
-Downloader = DOORDownloader
+Downloader = DOORDownloader

src/door/data_sources/chirps/chirps_downloader.py

@@ -6,7 +6,7 @@
 import datetime as dt
 import requests
 
-from ...base_downloaders import URLDownloader
+from ...base_downloaders import FTPDownloader
 
 from d3tools import timestepping as ts
 from d3tools.timestepping.timestep import TimeStep
@@ -15,15 +15,15 @@
 
 from ...utils.io import decompress_gz
 
-class CHIRPSDownloader(URLDownloader):
+class CHIRPSDownloader(FTPDownloader):
     source = "CHIRPS"
     name = "CHIRPS_downloader"
 
     default_options = {
         'get_prelim' : True, # if True, will also download preliminary data if available
     }
 
-    host = 'ftp://ftp.chc.ucsb.edu'
+    host = 'ftp.chc.ucsb.edu'
     auth = ('anonymous', '')
 
     homev2 = "/pub/org/chc/products/CHIRPS-2.0/"
@@ -76,7 +76,7 @@ class CHIRPSDownloader(URLDownloader):
 
     def __init__(self, product: str) -> None:
         self.set_product(product)
-        super().__init__(self.url_blank, protocol = 'ftp', host = self.host)
+        super().__init__(self.host, protocol = 'ftp')
 
     def set_product(self, product: str) -> None:
         self.product = product
@@ -100,28 +100,27 @@ def get_last_published_ts(self, prelim = None, product = None, **kwargs) -> ts.T
             product = self.product
 
         ts_per_year = self.available_products[product]["ts_per_year"]
-        url = self.available_products[product]["url"] if not prelim else self.available_products[product]["prelim_url"]
+        blank_path = self.available_products[product]["url"] if not prelim else self.available_products[product]["prelim_url"]
 
         if ts_per_year == 365:
             TimeStep = ts.Day
         else:
             TimeStep = FixedNTimeStep.get_subclass(ts_per_year)
 
         current_timestep = TimeStep.from_date(dt.datetime.now())
-        while True:
-            if "pentad_of_month" in url:
+        if 'lim' in kwargs:
+            start = TimeStep.from_date(kwargs['lim'])
+        else:
+            start = TimeStep.from_date(dt.datetime(2000, 1, 1))
+        
+        while current_timestep >= start:
+            if "pentad_of_month" in blank_path:
                 pentad_of_month = 6 if ts_per_year == 12 else current_timestep.dekad_of_month*2
-                current_url = url.format(timestep = current_timestep, pentad_of_month = pentad_of_month)
+                current_url = blank_path.format(timestep = current_timestep, pentad_of_month = pentad_of_month)
             else:   
-                current_url = url.format(timestep = current_timestep)
+                current_url = blank_path.format(timestep = current_timestep)
 
-            # send a request to the url
-            from ftpretty import ftpretty as ftp
-            client = ftp(self.host.replace('ftp://', ''), self.auth[0], self.auth[1])
-            list = client.list(current_url)
-
-            # if the request is successful, the last published timestep is the current timestep
-            if len(list) > 0:
+            if self.check_data(current_url):
                 return current_timestep
 
             # if the request is not successful, move to the previous timestep
@@ -140,7 +139,7 @@ def _get_data_ts(self,
         tmp_filename = f'{tmp_filename_raw}.tif.gz' if self.url_blank.endswith('.gz') else f'{tmp_filename_raw}.tif'
         tmp_destination = os.path.join(tmp_path, tmp_filename)
         if "pentad_of_month" not in self.url_blank:
-            success = self.download(tmp_destination, min_size = 200, missing_action = 'ignore', timestep = timestep, auth = self.auth)
+            success = self.download(self.url_blank, tmp_destination, min_size = 200, missing_action = 'ignore', timestep = timestep, auth = self.auth)
         else:
             if isinstance(timestep, ts.Month):
                 pentads = [1, 2, 3, 4, 5, 6]
@@ -151,7 +150,7 @@ def _get_data_ts(self,
             for pentad in pentads:
                 tmp_filename = f'{tmp_filename_raw}_{pentad}.tif'
                 tmp_destination = os.path.join(tmp_path, tmp_filename)
-                success = self.download(tmp_destination, min_size = 200, missing_action = 'ignore', timestep = timestep, pentad_of_month = pentad, auth = self.auth)
+                success = self.download(self.url_blank, tmp_destination, min_size = 200, missing_action = 'ignore', timestep = timestep, pentad_of_month = pentad, auth = self.auth)
                 tmp_filenames.append(tmp_filename)
                 if not success:
                     break
@@ -171,7 +170,7 @@ def _get_data_ts(self,
             if "pentad_of_month" not in self.url_prelim_blank:
                 tmp_filename = f'{tmp_filename_raw}.tif.gz' if self.url_prelim_blank.endswith('.gz') else f'{tmp_filename_raw}.tif'
                 tmp_destination = os.path.join(tmp_path, tmp_filename)
-                success = self.download(tmp_destination, min_size = 200, missing_action = 'ignore', timestep = timestep, prelim = True, auth = self.auth)
+                success = self.download(self.url_prelim_blank, tmp_destination, min_size = 200, missing_action = 'ignore', timestep = timestep, prelim = True, auth = self.auth)
 
             else:
                 if isinstance(timestep, ts.Month):
@@ -183,7 +182,7 @@ def _get_data_ts(self,
                 for pentad in pentads:
                     tmp_filename = f'{tmp_filename_raw}_{pentad}.tif'
                     tmp_destination = os.path.join(tmp_path, tmp_filename)
-                    success = self.download(tmp_destination, min_size = 200, missing_action = 'ignore', timestep = timestep, pentad_of_month = pentad, prelim = True, auth = self.auth)
+                    success = self.download(self.url_prelim_blank, tmp_destination, min_size = 200, missing_action = 'ignore', timestep = timestep, pentad_of_month = pentad, prelim = True, auth = self.auth)
                     tmp_filenames.append(tmp_filename)
                     if not success:
                         break
@@ -215,14 +214,4 @@ def _get_data_ts(self,
             if isprelim:
                 cropped.attrs['PRELIMINARY'] = 'True'
 
-            yield cropped, {}
-    
-    def format_url(self, prelim = False, **kwargs):
-        """
-        Format the url for the download
-        """
-        if prelim:
-            url = self.url_prelim_blank.format(**kwargs)
-        else:
-            url = self.url_blank.format(**kwargs)
-        return url
+            yield cropped, {}