import logging
import os
import re
from builtins import object
from pathlib import Path
import requests
from requests import Session
from requests.adapters import HTTPAdapter
from urllib3.util import Retry
from threeML.config.config import threeML_config
from threeML.io.file_utils import (
file_existing_and_readable,
path_exists_and_is_directory,
sanitize_filename,
)
from threeML.utils.progress_bar import tqdm
log = logging.getLogger(__name__)
[docs]
class RemoteDirectoryNotFound(IOError):
pass
[docs]
class HTTPError(IOError):
pass
[docs]
class ApacheDirectory(object):
"""Allows to interact with a directory listing like the one returned by an
Apache server."""
def __init__(self, url):
self._request_result = requests.get(url)
# Make sure the request was ok
if not self._request_result.ok:
if self._request_result.reason == "Not Found":
raise RemoteDirectoryNotFound(
"Remote directory %s does not exist" % url
)
else:
raise HTTPError(
"HTTP request failed with reason: %s" % self._request_result.reason
)
self._text = self._request_result.text
# Get the listing of files and directories
self._entries = self._get_directory_entries()
# Now split directories and files
self._files = []
self._directories = []
for entry in self._entries:
if entry[1] == "FILE":
self._files.append(entry[0])
else:
self._directories.append(entry[0])
def _get_directory_entries(self):
"""List files and directories listed in the listing.
:return: a list of tuples (entry name, type (DIR or FILE))
"""
# Get the files listed in the directory
# A line in an Apache listing is like this:
# <img src="/icons/unknown.gif" alt="[ ]">
# <a href="glg_cspec_b0_bn100101988_v02.rsp">glg_cspec_b0_bn100101988_v02.rsp
# </a>
# 16-Nov-2012 15:14 96K
regexp = re.compile(r"<img src=.+ alt=(.+)>\s?<a href=.+>(.+)</a>.+")
# Apache puts files in a <pre></pre> tag, so lines are ended simply with \n
lines = self._text.split("\n")
# Now loop over the lines and extract the file name
entries = []
for line in lines:
token = re.match(regexp, line)
if token is not None:
# This line contains a file or a directory
type_token, filename_token = token.groups()
# Figure out if this is a directory or a file. A directory has a
# alt="[DIR]" attribute in the
# <img> tag, a file has a alt="[ ]" or other things (if a known type)
if type_token.upper().find("DIR") >= 0:
entry_type = "DIR"
else:
entry_type = "FILE"
# Append entry
entries.append((filename_token, entry_type))
return entries
@property
def files(self):
return self._files
@property
def directories(self):
return self._directories
[docs]
def download(
self,
remote_filename,
destination_path: str,
new_filename=None,
progress=True,
compress=False,
timeout=10,
):
assert (
remote_filename in self.files
), "File %s is not contained in this directory (%s)" % (
remote_filename,
self._request_result.url,
)
destination_path: Path = sanitize_filename(destination_path, abspath=True)
assert path_exists_and_is_directory(destination_path), (
f"Provided destination {destination_path} does not exist or "
"is not a directory"
)
# If no filename is specified, use the same name that the file has on the remote
# server
if new_filename is None:
new_filename: str = remote_filename.split("/")[-1]
# Get the fully qualified path for the remote and the local file
remote_path: str = self._request_result.url + remote_filename
local_path: Path = destination_path / new_filename
# Ask the server for the file, but do not download it just yet
# (stream=True will get the HTTP header but nothing else)
# Use stream=True for two reasons:
# * so that the file is not downloaded all in memory before being written to the
# disk
# * so that we can report progress is requested
with Session() as session: # use a Session context manager to allow retries
# using requests
retries = Retry(total=3, backoff_factor=0.1)
session.mount("https://", HTTPAdapter(max_retries=retries))
session.mount("http://", HTTPAdapter(max_retries=retries))
this_request = session.get(remote_path, stream=True, timeout=timeout)
# Figure out the size of the file
file_size = int(this_request.headers["Content-Length"])
log.debug(f"downloading {remote_filename} of size {file_size}")
# Now check if we really need to download this file
if compress:
# Add a .gz at the end of the file path
log.debug(f"file {remote_filename} will be downloaded and compressed")
local_path: Path = Path(f"{local_path}.gz")
if file_existing_and_readable(local_path):
local_size = os.path.getsize(local_path)
if local_size == file_size or compress:
# if the compressed file already exists
# it will have a smaller size
# No need to download it again
log.info(f"file {remote_filename} is already downloaded!")
return local_path
if local_path.is_file():
first_byte = os.path.getsize(local_path)
else:
first_byte = 0
# Chunk size shouldn't bee too small otherwise we are causing a bottleneck
# in the download speed
chunk_size = 1024 * 10
# If the user wants to compress the file, use gzip, otherwise the normal
# opener
if compress:
import gzip
opener = gzip.open
else:
opener = open
if threeML_config["interface"]["progress_bars"]:
# Set a title for the progress bar
bar_title = "Downloading %s" % new_filename
bar = tqdm(
initial=first_byte,
unit_scale=True,
unit_divisor=1024,
unit="B",
total=int(this_request.headers["Content-Length"]),
desc=bar_title,
)
with opener(local_path, "wb") as f:
for chunk in this_request.iter_content(chunk_size=chunk_size):
if chunk: # filter out keep-alive new chunks
f.write(chunk)
bar.update(len(chunk))
bar.close()
else:
with opener(local_path, "wb") as f:
for chunk in this_request.iter_content(chunk_size=chunk_size):
if chunk: # filter out keep-alive new chunks
f.write(chunk)
return local_path
[docs]
def download_all_files(self, destination_path, progress=True, pattern=None):
"""Download all files in the current directory.
:param destination_path: the path for the destination directory
in the local file system
:param progress: (True or False) whether to display progress or
not
:param pattern: (default: None) If not None, only files matching
this pattern (a regular expression) will be downloaded
:return: list of the downloaded files as absolute paths in the
local file system
"""
local_files = []
for file in self.files:
if pattern is not None:
if re.match(pattern, os.path.basename(file)) is None:
continue
this_local_file = self.download(file, destination_path, progress=progress)
local_files.append(this_local_file)
return local_files