diff --git a/pooch/__init__.py b/pooch/__init__.py index 826cb57d1..bef6053a6 100644 --- a/pooch/__init__.py +++ b/pooch/__init__.py @@ -19,7 +19,7 @@ from .processors import Unzip, Untar, Decompress # This file is generated automatically by setuptools_scm -from . import _version # type: ignore +from . import _version # type: ignore[attr-defined] # Add a "v" to the version number diff --git a/pooch/downloaders.py b/pooch/downloaders.py index 561640fb4..58d3fbf8a 100644 --- a/pooch/downloaders.py +++ b/pooch/downloaders.py @@ -7,25 +7,33 @@ """ The classes that actually handle the downloads. """ + import os import sys import ftplib +from io import BufferedRandom +from abc import abstractmethod import warnings +from typing import Union, TYPE_CHECKING, cast, Optional, Any from .utils import parse_url +from .typing import Downloader, ProgressBar, PathType + +if TYPE_CHECKING: + from .core import Pooch # Mypy doesn't like assigning None like this. # Can just use a guard variable try: from tqdm import tqdm except ImportError: - tqdm = None # type: ignore + tqdm = None # type: ignore[misc, assignment] try: import paramiko except ImportError: - paramiko = None # type: ignore + paramiko = None # type: ignore[misc, assignment] # Set the default timeout in seconds so it can be configured in a pinch for the @@ -33,8 +41,11 @@ # See https://github.com/fatiando/pooch/issues/409 DEFAULT_TIMEOUT = 30 +if TYPE_CHECKING: + ProgressBarArg = Union[ProgressBar, bool, tqdm] + -def choose_downloader(url, progressbar=False): +def choose_downloader(url: str, progressbar: "ProgressBarArg" = False) -> Downloader: """ Choose the appropriate downloader for the given URL based on the protocol. @@ -166,7 +177,9 @@ class HTTPDownloader: # pylint: disable=too-few-public-methods """ - def __init__(self, progressbar=False, chunk_size=1024, **kwargs): + def __init__( + self, progressbar: "ProgressBarArg" = False, chunk_size: int = 1024, **kwargs + ) -> None: self.kwargs = kwargs self.progressbar = progressbar self.chunk_size = chunk_size @@ -174,8 +187,8 @@ def __init__(self, progressbar=False, chunk_size=1024, **kwargs): raise ValueError("Missing package 'tqdm' required for progress bars.") def __call__( - self, url, output_file, pooch, check_only=False - ): # pylint: disable=R0914 + self, url: str, output_file: PathType, pooch: "Pooch", check_only: bool = False + ) -> Union[None, bool]: # pylint: disable=R0914 """ Download the given URL over HTTP to the given output file. @@ -216,8 +229,10 @@ def __call__( ispath = not hasattr(output_file, "write") if ispath: # pylint: disable=consider-using-with - output_file = open(output_file, "w+b") + output_file_writer = open(output_file, "w+b") # pylint: enable=consider-using-with + else: + output_file_writer = cast(BufferedRandom, output_file) try: response = requests.get(url, timeout=timeout, **kwargs) response.raise_for_status() @@ -228,7 +243,7 @@ def __call__( # always full unicode support # (see https://github.com/tqdm/tqdm/issues/454) use_ascii = bool(sys.platform == "win32") - progress = tqdm( + progress: ProgressBar = tqdm( total=total, ncols=79, ascii=use_ascii, @@ -237,12 +252,12 @@ def __call__( leave=True, ) elif self.progressbar: - progress = self.progressbar + progress: ProgressBar = self.progressbar # type: ignore[no-redef] progress.total = total for chunk in content: if chunk: - output_file.write(chunk) - output_file.flush() + output_file_writer.write(chunk) + output_file_writer.flush() if self.progressbar: # Use the chunk size here because chunk may be much # larger if the data are decompressed by requests after @@ -258,7 +273,7 @@ def __call__( progress.close() finally: if ispath: - output_file.close() + output_file_writer.close() return None @@ -286,7 +301,7 @@ class FTPDownloader: # pylint: disable=too-few-public-methods to indicate no password is required. account : str Some servers also require an "account" name for authentication. - timeout : int + timeout : Optional[int] Timeout in seconds for ftp socket operations, use None to mean no timeout. progressbar : bool @@ -301,14 +316,14 @@ class FTPDownloader: # pylint: disable=too-few-public-methods def __init__( self, - port=21, - username="anonymous", - password="", - account="", - timeout=None, - progressbar=False, - chunk_size=1024, - ): + port: int = 21, + username: str = "anonymous", + password: str = "", + account: str = "", + timeout: Optional[int] = None, + progressbar: bool = False, + chunk_size: int = 1024, + ) -> None: self.port = port self.username = username self.password = password @@ -319,7 +334,9 @@ def __init__( if self.progressbar is True and tqdm is None: raise ValueError("Missing package 'tqdm' required for progress bars.") - def __call__(self, url, output_file, pooch, check_only=False): + def __call__( + self, url: str, output_file: PathType, pooch: "Pooch", check_only: bool = False + ) -> Union[None, bool]: """ Download the given URL over FTP to the given output file. @@ -359,8 +376,10 @@ def __call__(self, url, output_file, pooch, check_only=False): ispath = not hasattr(output_file, "write") if ispath: # pylint: disable=consider-using-with - output_file = open(output_file, "w+b") + output_file_writer = open(output_file, "w+b") # pylint: enable=consider-using-with + else: + output_file_writer = cast(BufferedRandom, output_file) try: ftp.login(user=self.username, passwd=self.password, acct=self.account) command = f"RETR {parsed_url['path']}" @@ -369,8 +388,9 @@ def __call__(self, url, output_file, pooch, check_only=False): # get the file size. See: https://stackoverflow.com/a/22093848 ftp.voidcmd("TYPE I") use_ascii = bool(sys.platform == "win32") + file_size = ftp.size(parsed_url["path"]) progress = tqdm( - total=int(ftp.size(parsed_url["path"])), + total=int(file_size) if file_size is not None else 0, ncols=79, ascii=use_ascii, unit="B", @@ -382,15 +402,17 @@ def __call__(self, url, output_file, pooch, check_only=False): def callback(data): "Update the progress bar and write to output" progress.update(len(data)) - output_file.write(data) + output_file_writer.write(data) ftp.retrbinary(command, callback, blocksize=self.chunk_size) else: - ftp.retrbinary(command, output_file.write, blocksize=self.chunk_size) + ftp.retrbinary( + command, output_file_writer.write, blocksize=self.chunk_size + ) finally: ftp.quit() if ispath: - output_file.close() + output_file_writer.close() return None @@ -417,7 +439,7 @@ class SFTPDownloader: # pylint: disable=too-few-public-methods Password used to login to the server. Only needed if the server requires authentication (i.e., no anonymous SFTP). Use the empty string to indicate no password is required. - timeout : int + timeout : Optional[int] Timeout in seconds for sftp socket operations, use None to mean no timeout. progressbar : bool or an arbitrary progress bar object @@ -429,13 +451,13 @@ class SFTPDownloader: # pylint: disable=too-few-public-methods def __init__( self, - port=22, - username="anonymous", - password="", - account="", - timeout=None, - progressbar=False, - ): + port: int = 22, + username: str = "anonymous", + password: str = "", + account: str = "", + timeout: Optional[int] = None, + progressbar: "ProgressBarArg" = False, + ) -> None: self.port = port self.username = username self.password = password @@ -445,7 +467,7 @@ def __init__( # Collect errors and raise only once so that both missing packages are # captured. Otherwise, the user is only warned of one of them at a # time (and we can't test properly when they are both missing). - errors = [] + errors: list[str] = [] if self.progressbar and tqdm is None: errors.append("Missing package 'tqdm' required for progress bars.") if paramiko is None: @@ -453,7 +475,7 @@ def __init__( if errors: raise ValueError(" ".join(errors)) - def __call__(self, url, output_file, pooch): + def __call__(self, url: str, output_file: str, pooch: "Pooch") -> None: """ Download the given URL over SFTP to the given output file. @@ -476,12 +498,16 @@ def __call__(self, url, output_file, pooch): try: connection.connect(username=self.username, password=self.password) sftp = paramiko.SFTPClient.from_transport(connection) - sftp.get_channel().settimeout = self.timeout + + if sftp is None: + raise ValueError + + sftp.get_channel().settimeout = self.timeout # type: ignore[method-assign, union-attr] if self.progressbar: - size = int(sftp.stat(parsed_url["path"]).st_size) + size = sftp.stat(parsed_url["path"]).st_size use_ascii = bool(sys.platform == "win32") progress = tqdm( - total=size, + total=int(size) if size is not None else 0, ncols=79, ascii=use_ascii, unit="B", @@ -587,12 +613,14 @@ class DOIDownloader: # pylint: disable=too-few-public-methods """ - def __init__(self, progressbar=False, chunk_size=1024, **kwargs): + def __init__( + self, progressbar: "ProgressBarArg" = False, chunk_size: int = 1024, **kwargs + ) -> None: self.kwargs = kwargs self.progressbar = progressbar self.chunk_size = chunk_size - def __call__(self, url, output_file, pooch): + def __call__(self, url: str, output_file: PathType, pooch: "Pooch") -> None: """ Download the given DOI URL over HTTP to the given output file. @@ -629,7 +657,7 @@ def __call__(self, url, output_file, pooch): downloader(download_url, output_file, pooch) -def doi_to_url(doi): +def doi_to_url(doi: str) -> str: """ Follow a DOI link to resolve the URL of the archive. @@ -657,7 +685,7 @@ def doi_to_url(doi): return url -def doi_to_repository(doi): +def doi_to_repository(doi: str) -> "DataRepository": """ Instantiate a data repository instance from a given DOI. @@ -712,7 +740,10 @@ def doi_to_repository(doi): class DataRepository: # pylint: disable=too-few-public-methods, missing-class-docstring @classmethod - def initialize(cls, doi, archive_url): # pylint: disable=unused-argument + @abstractmethod + def initialize( + cls, doi: str, archive_url: str + ) -> Union[None, "DataRepository"]: # pylint: disable=unused-argument """ Initialize the data repository if the given URL points to a corresponding repository. @@ -732,7 +763,8 @@ def initialize(cls, doi, archive_url): # pylint: disable=unused-argument return None # pragma: no cover - def download_url(self, file_name): + @abstractmethod + def download_url(self, file_name: str) -> str: """ Use the repository API to get the download URL for a file given the archive URL. @@ -750,7 +782,8 @@ def download_url(self, file_name): raise NotImplementedError # pragma: no cover - def populate_registry(self, pooch): + @abstractmethod + def populate_registry(self, pooch: "Pooch") -> None: """ Populate the registry using the data repository's API @@ -766,14 +799,14 @@ def populate_registry(self, pooch): class ZenodoRepository(DataRepository): # pylint: disable=missing-class-docstring base_api_url = "https://zenodo.org/api/records" - def __init__(self, doi, archive_url): + def __init__(self, doi: str, archive_url: str) -> None: self.archive_url = archive_url self.doi = doi self._api_response = None - self._api_version = None + self._api_version: Union[None, str] = None @classmethod - def initialize(cls, doi, archive_url): + def initialize(cls, doi: str, archive_url: str) -> Union[None, "ZenodoRepository"]: """ Initialize the data repository if the given URL points to a corresponding repository. @@ -799,7 +832,7 @@ def initialize(cls, doi, archive_url): return cls(doi, archive_url) @property - def api_response(self): + def api_response(self) -> Any: """Cached API response from Zenodo""" if self._api_response is None: # Lazy import requests to speed up import time @@ -814,7 +847,7 @@ def api_response(self): return self._api_response @property - def api_version(self): + def api_version(self) -> str: """ Version of the Zenodo API we are interacting with @@ -845,7 +878,7 @@ def api_version(self): ) return self._api_version - def download_url(self, file_name): + def download_url(self, file_name: str) -> str: """ Use the repository API to get the download URL for a file given the archive URL. @@ -873,7 +906,7 @@ def download_url(self, file_name): if self.api_version == "legacy": files = {item["key"]: item for item in self.api_response["files"]} else: - files = [item["filename"] for item in self.api_response["files"]] + files = {item["filename"]: None for item in self.api_response["files"]} # Check if file exists in the repository if file_name not in files: raise ValueError( @@ -890,7 +923,7 @@ def download_url(self, file_name): ) return download_url - def populate_registry(self, pooch): + def populate_registry(self, pooch: "Pooch") -> None: """ Populate the registry using the data repository's API @@ -917,13 +950,15 @@ def populate_registry(self, pooch): class FigshareRepository(DataRepository): # pylint: disable=missing-class-docstring - def __init__(self, doi, archive_url): + def __init__(self, doi: str, archive_url: str) -> None: self.archive_url = archive_url self.doi = doi self._api_response = None @classmethod - def initialize(cls, doi, archive_url): + def initialize( + cls, doi: str, archive_url: str + ) -> Union[None, "FigshareRepository"]: """ Initialize the data repository if the given URL points to a corresponding repository. @@ -948,7 +983,7 @@ def initialize(cls, doi, archive_url): return cls(doi, archive_url) - def _parse_version_from_doi(self): + def _parse_version_from_doi(self) -> Union[None, int]: """ Parse version from the doi @@ -965,7 +1000,7 @@ def _parse_version_from_doi(self): return version @property - def api_response(self): + def api_response(self) -> Any: """Cached API response from Figshare""" if self._api_response is None: # Lazy import requests to speed up import time @@ -1007,7 +1042,7 @@ def api_response(self): return self._api_response - def download_url(self, file_name): + def download_url(self, file_name: str) -> str: """ Use the repository API to get the download URL for a file given the archive URL. @@ -1030,7 +1065,7 @@ def download_url(self, file_name): download_url = files[file_name]["download_url"] return download_url - def populate_registry(self, pooch): + def populate_registry(self, pooch: "Pooch") -> None: """ Populate the registry using the data repository's API @@ -1045,13 +1080,15 @@ def populate_registry(self, pooch): class DataverseRepository(DataRepository): # pylint: disable=missing-class-docstring - def __init__(self, doi, archive_url): + def __init__(self, doi: str, archive_url: str) -> None: self.archive_url = archive_url self.doi = doi self._api_response = None @classmethod - def initialize(cls, doi, archive_url): + def initialize( + cls, doi: str, archive_url: str + ) -> Union[None, "DataverseRepository"]: """ Initialize the data repository if the given URL points to a corresponding repository. @@ -1081,7 +1118,7 @@ def initialize(cls, doi, archive_url): return repository @classmethod - def _get_api_response(cls, doi, archive_url): + def _get_api_response(cls, doi: str, archive_url: str) -> Any: """ Perform the actual API request @@ -1100,7 +1137,7 @@ def _get_api_response(cls, doi, archive_url): return response @property - def api_response(self): + def api_response(self) -> Any: """Cached API response from a DataVerse instance""" if self._api_response is None: @@ -1111,12 +1148,12 @@ def api_response(self): return self._api_response @api_response.setter - def api_response(self, response): + def api_response(self, response: Any) -> Any: """Update the cached API response""" self._api_response = response - def download_url(self, file_name): + def download_url(self, file_name: str) -> str: """ Use the repository API to get the download URL for a file given the archive URL. @@ -1149,7 +1186,7 @@ def download_url(self, file_name): ) return download_url - def populate_registry(self, pooch): + def populate_registry(self, pooch: "Pooch") -> None: """ Populate the registry using the data repository's API diff --git a/pooch/hashes.py b/pooch/hashes.py index ebac68b97..6662e6f2e 100644 --- a/pooch/hashes.py +++ b/pooch/hashes.py @@ -7,9 +7,12 @@ """ Calculating and checking file hashes. """ + import hashlib import functools from pathlib import Path +from typing import Optional +from .typing import PathType # From the docs: https://docs.python.org/3/library/hashlib.html#hashlib.new # The named constructors are much faster than new() and should be @@ -40,7 +43,7 @@ pass -def file_hash(fname, alg="sha256"): +def file_hash(fname: PathType, alg: str = "sha256") -> str: """ Calculate the hash of a given file. @@ -48,8 +51,8 @@ def file_hash(fname, alg="sha256"): Parameters ---------- - fname : str - The name of the file. + fname : str or PathLike + The path to the file. alg : str The type of the hashing algorithm @@ -87,7 +90,7 @@ def file_hash(fname, alg="sha256"): return hasher.hexdigest() -def hash_algorithm(hash_string): +def hash_algorithm(hash_string: str) -> str: """ Parse the name of the hash method from the hash string. @@ -134,7 +137,12 @@ def hash_algorithm(hash_string): return algorithm.lower() -def hash_matches(fname, known_hash, strict=False, source=None): +def hash_matches( + fname: PathType, + known_hash: Optional[str], + strict: bool = False, + source: Optional[str] = None, +) -> bool: """ Check if the hash of a file matches a known hash. @@ -147,13 +155,13 @@ def hash_matches(fname, known_hash, strict=False, source=None): ---------- fname : str or PathLike The path to the file. - known_hash : str + known_hash : Optional[str] The known hash. Optionally, prepend ``alg:`` to the hash to specify the hashing algorithm. Default is SHA256. strict : bool If True, will raise a :class:`ValueError` if the hash does not match informing the user that the file may be corrupted. - source : str + source : Optional[str] The source of the downloaded file (name or URL, for example). Will be used in the error message if *strict* is True. Has no other use other than reporting to the user where the file came from in case of hash @@ -182,7 +190,7 @@ def hash_matches(fname, known_hash, strict=False, source=None): return matches -def make_registry(directory, output, recursive=True): +def make_registry(directory: str, output: str, recursive: bool = True) -> None: """ Make a registry of files and hashes for the given directory. @@ -201,19 +209,19 @@ def make_registry(directory, output, recursive=True): *directory*. """ - directory = Path(directory) + directory_path = Path(directory) if recursive: pattern = "**/*" else: pattern = "*" files = sorted( - str(path.relative_to(directory)) - for path in directory.glob(pattern) + str(path.relative_to(directory_path)) + for path in directory_path.glob(pattern) if path.is_file() ) - hashes = [file_hash(str(directory / fname)) for fname in files] + hashes = [file_hash(str(directory_path / fname)) for fname in files] with open(output, "w", encoding="utf-8") as outfile: for fname, fhash in zip(files, hashes): diff --git a/pooch/processors.py b/pooch/processors.py index dfeebb387..b5fd9826f 100644 --- a/pooch/processors.py +++ b/pooch/processors.py @@ -19,8 +19,12 @@ from zipfile import ZipFile from tarfile import TarFile +from typing import Any, Union, TYPE_CHECKING from .utils import get_logger +if TYPE_CHECKING: + from .core import Pooch + class ExtractorProcessor(abc.ABC): # pylint: disable=too-few-public-methods """ @@ -46,13 +50,15 @@ class ExtractorProcessor(abc.ABC): # pylint: disable=too-few-public-methods """ - def __init__(self, members=None, extract_dir=None): + def __init__( + self, members: Union[None, list] = None, extract_dir: Union[None, str] = None + ) -> None: self.members = members self.extract_dir = extract_dir @property @abc.abstractmethod - def suffix(self): + def suffix(self) -> str: """ String appended to unpacked archive folder name. Only used if extract_dir is None. @@ -60,21 +66,21 @@ def suffix(self): """ @abc.abstractmethod - def _all_members(self, fname): + def _all_members(self, fname: str) -> list[str]: """ Return all the members in the archive. MUST BE IMPLEMENTED BY CHILD CLASSES. """ @abc.abstractmethod - def _extract_file(self, fname, extract_dir): + def _extract_file(self, fname: str, extract_dir: str) -> None: """ This method receives an argument for the archive to extract and the destination path. MUST BE IMPLEMENTED BY CHILD CLASSES. """ - def __call__(self, fname, action, pooch): + def __call__(self, fname: str, action: str, pooch: "Pooch") -> list[str]: """ Extract all files from the given archive. @@ -166,19 +172,19 @@ class Unzip(ExtractorProcessor): # pylint: disable=too-few-public-methods """ @property - def suffix(self): + def suffix(self) -> str: """ String appended to unpacked archive folder name. Only used if extract_dir is None. """ return ".unzip" - def _all_members(self, fname): + def _all_members(self, fname: str) -> list[str]: """Return all members from a given archive.""" with ZipFile(fname, "r") as zip_file: return zip_file.namelist() - def _extract_file(self, fname, extract_dir): + def _extract_file(self, fname: str, extract_dir: str) -> None: """ This method receives an argument for the archive to extract and the destination path. @@ -238,19 +244,19 @@ class Untar(ExtractorProcessor): # pylint: disable=too-few-public-methods """ @property - def suffix(self): + def suffix(self) -> str: """ String appended to unpacked archive folder name. Only used if extract_dir is None. """ return ".untar" - def _all_members(self, fname): + def _all_members(self, fname: str) -> list[str]: """Return all members from a given archive.""" with TarFile.open(fname, "r") as tar_file: return [info.name for info in tar_file.getmembers()] - def _extract_file(self, fname, extract_dir): + def _extract_file(self, fname: str, extract_dir: str) -> None: """ This method receives an argument for the archive to extract and the destination path. @@ -262,7 +268,10 @@ def _extract_file(self, fname, extract_dir): "Untarring contents of '%s' to '%s'", fname, extract_dir ) # Unpack all files from the archive into our new folder - tar_file.extractall(path=extract_dir, **filter_kwarg) + tar_file.extractall( + path=extract_dir, + **filter_kwarg, # type: ignore[arg-type] + ) else: for member in self.members: get_logger().info( @@ -285,7 +294,9 @@ def _extract_file(self, fname, extract_dir): ] # Extract the data file from within the archive tar_file.extractall( - members=subdir_members, path=extract_dir, **filter_kwarg + members=subdir_members, + path=extract_dir, + **filter_kwarg, # type: ignore[arg-type] ) @@ -336,11 +347,11 @@ class Decompress: # pylint: disable=too-few-public-methods modules = {"auto": None, "lzma": lzma, "xz": lzma, "gzip": gzip, "bzip2": bz2} extensions = {".xz": "lzma", ".gz": "gzip", ".bz2": "bzip2"} - def __init__(self, method="auto", name=None): + def __init__(self, method: str = "auto", name: Union[None, str] = None) -> None: self.method = method self.name = name - def __call__(self, fname, action, pooch): + def __call__(self, fname: str, action: str, pooch: "Pooch") -> str: """ Decompress the given file. @@ -384,7 +395,7 @@ class attribute. shutil.copyfileobj(compressed, output) return decompressed - def _compression_module(self, fname): + def _compression_module(self, fname: str) -> Any: """ Get the Python module compatible with fname and the chosen method. diff --git a/pooch/py.typed b/pooch/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/pooch/tests/test_downloaders.py b/pooch/tests/test_downloaders.py index 9b089ede0..b5ab76d5d 100644 --- a/pooch/tests/test_downloaders.py +++ b/pooch/tests/test_downloaders.py @@ -19,12 +19,12 @@ try: import tqdm except ImportError: - tqdm = None # type: ignore + tqdm = None # type: ignore[assignment] try: import paramiko except ImportError: - paramiko = None # type: ignore + paramiko = None # type: ignore[assignment] from .. import Pooch from ..downloaders import ( diff --git a/pooch/typing/__init__.py b/pooch/typing/__init__.py index c683490ce..1d7c7b612 100644 --- a/pooch/typing/__init__.py +++ b/pooch/typing/__init__.py @@ -60,6 +60,23 @@ def __call__( # noqa: E704 ) -> Any: ... +class ProgressBar(Protocol): + """ + Class used to define the type definition for a progress bar. + """ + + total: int + + # pylint: disable=too-few-public-methods + def update( + self, n: Union[float, None] + ) -> Union[None, bool]: ... # noqa: E704, C0116 + + def reset(self) -> None: ... # noqa: E704, C0116 + + def close(self) -> None: ... # noqa: E704, C0116 + + class ParsedURL(TypedDict): """ Type for a dictionary generated after parsing a URL. diff --git a/pooch/utils.py b/pooch/utils.py index c0c74a426..3e710fd50 100644 --- a/pooch/utils.py +++ b/pooch/utils.py @@ -303,7 +303,7 @@ def temporary_file(path: Optional[PathType] = None) -> Generator[str, None, None The path to the temporary file. """ - tmp = tempfile.NamedTemporaryFile(delete=False, dir=path) # type: ignore + tmp = tempfile.NamedTemporaryFile(delete=False, dir=path) # type: ignore[type-var] # Close the temp file so that it can be opened elsewhere tmp.close() try: