Repository: cogent3/cogent3
Pull request #1069 (2022-06-29 05:45:17)

ENH: open_ now handles urls too
[NEW] this enables files to be downloaded from the web using the
    convenience functions,
    e.g. load_aligned_seqs("https://path/to/somefile.fasta")
    will now work.

[NEW] open_url() now supports "file://path/to/local/file". Mainly for
    testing purposes.

[CHANGED] tests of open_url() now only use 1 remote file to reduce
    test time.


Patch:
diff --git a/src/cogent3/parse/fasta.py b/src/cogent3/parse/fasta.py
index 81edf298a..6ea76124a 100644
--- a/src/cogent3/parse/fasta.py
+++ b/src/cogent3/parse/fasta.py
@@ -61,16 +61,15 @@ def MinimalFastaParser(
     try:
         infile = open_(infile)
         close_at_end = True
-    except (TypeError, AttributeError):
+    except (ValueError, TypeError, AttributeError):
         close_at_end = False
 
     for rec in finder(infile):
         # first line must be a label line
-        if not rec[0][0] in label_characters:
+        if rec[0][0] not in label_characters:
             if strict:
                 raise RecordError(f"Found Fasta record without label line: {rec}")
             continue
-        # record must have at least one sequence
         if len(rec) < 2:
             if strict:
                 raise RecordError(f"Found label line without sequences: {rec}")
diff --git a/src/cogent3/util/io.py b/src/cogent3/util/io.py
index bc84d69c4..2377cff8a 100644
--- a/src/cogent3/util/io.py
+++ b/src/cogent3/util/io.py
@@ -7,8 +7,9 @@ from io import TextIOWrapper
 from os import path as os_path
 from os import remove
 from pathlib import Path
+from re import compile
 from tempfile import mkdtemp
-from typing import Union
+from typing import IO, Callable, Union
 from urllib.parse import ParseResult, urlparse
 from urllib.request import urlopen
 from zipfile import ZipFile
@@ -27,8 +28,33 @@ __maintainer__ = "Gavin Huttley"
 __email__ = "Gavin.Huttley@anu.edu.au"
 __status__ = "Production"
 
+# support prefixes for urls
+_urls = compile("^(http[s]*|file)")
 
-def open_zip(filename: Union[str, Path], mode="r", **kwargs):
+
+def _get_compression_open(
+    path: Union[str, Path] = None, compression: Union[str] = None
+) -> Union[Callable]:
+    """returns function for opening compression formats
+
+    Parameters
+    ----------
+    path
+        file path or url
+    compression
+        file compression suffix
+
+    Returns
+    -------
+    function for opening compressed files or None if unknown compression
+    """
+    assert path or compression
+    if compression is None:
+        _, compression = get_format_suffixes(path)
+    return {"gz": gzip_open, "bz2": bzip_open, "zip": open_zip}.get(compression, None)
+
+
+def open_zip(filename: Union[str, Path], mode: str = "r", **kwargs) -> IO:
     """open a single member zip-compressed file
 
     Note
@@ -62,28 +88,44 @@ def open_zip(filename: Union[str, Path], mode="r", **kwargs):
         return TextIOWrapper(opened, encoding=encoding)
 
 
-def open_(filename: Union[str, Path], mode="rt", **kwargs):
-    """open that handles different compression"""
+def open_(filename: Union[str, Path], mode="rt", **kwargs) -> IO:
+    """open that handles different compression
+
+    Parameters
+    ----------
+    filename
+        path or url, if a url delegates processing to open_url
+    mode
+        standard file opening mode
+    kwargs
+        passed to open functions
+
+    Returns
+    -------
+    an object compatible with the file protocol
+    """
+    if not filename:
+        raise ValueError(f"{filename} not a valid file name or url")
+
+    if _urls.search(str(filename)):
+        return open_url(filename, mode=mode, **kwargs)
+
     mode = mode or "rt"
-    filename = Path(filename).expanduser().absolute()
-    op = {".gz": gzip_open, ".bz2": bzip_open, ".zip": open_zip}.get(
-        filename.suffix, open
-    )
+    op = _get_compression_open(filename) or open
 
     encoding = kwargs.pop("encoding", None)
     need_encoding = mode.startswith("r") and "b" not in mode
-    if need_encoding:
-        if "encoding" not in kwargs:
-            with op(filename, mode="rb") as infile:
-                data = infile.read(100)
+    if need_encoding and "encoding" not in kwargs:
+        with op(filename, mode="rb") as infile:
+            data = infile.read(100)
 
-            encoding = detect(data)
-            encoding = encoding["encoding"]
+        encoding = detect(data)
+        encoding = encoding["encoding"]
 
     return op(filename, mode, encoding=encoding, **kwargs)
 
 
-def open_url(url: Union[str, ParseResult], mode="r", **kwargs):
+def open_url(url: Union[str, ParseResult], mode="r", **kwargs) -> IO:
     """open a url
 
     Parameters
@@ -102,22 +144,28 @@ def open_url(url: Union[str, ParseResult], mode="r", **kwargs):
     If mode='b' or 'rb' (binary read), the function returns file object to read
     else returns TextIOWrapper to read text with specified encoding in the URL
     """
-    mode = mode or "r"
+    _, compression = get_format_suffixes(
+        getattr(url, "path", url)
+    )  # handling possibility of ParseResult
+    mode = "rb" if compression else mode or "r"
+
     url_parsed = url if isinstance(url, ParseResult) else urlparse(url)
 
     if "r" not in mode:
         raise ValueError("opening a url only allowed in read mode")
 
-    if url_parsed.scheme not in "https":
-        raise ValueError("URL scheme must be http or https.")
+    if not _urls.search(url_parsed.scheme):
+        raise ValueError(
+            f"URL scheme must be http, https or file, not {url_parsed.scheme}"
+        )
 
     response = urlopen(url_parsed.geturl(), timeout=10)
+    encoding = response.headers.get_content_charset()
+    if compression:
+        response = _get_compression_open(compression=compression)(response)
+        mode = "r"  # wrap it as text
 
-    return (
-        response
-        if "b" in mode
-        else TextIOWrapper(response, encoding=response.headers.get_content_charset())
-    )
+    return response if "b" in mode else TextIOWrapper(response, encoding=encoding)
 
 
 def _path_relative_to_zip_parent(zip_path, member_path):
@@ -224,7 +272,7 @@ class atomic_write:
 
         return self._file
 
-    def __enter__(self):
+    def __enter__(self) -> IO:
         return self._get_fileobj()
 
     def _close_rename_standard(self, src):


Test patch:
diff --git a/tests/test_util/test_io.py b/tests/test_util/test_io.py
index 31ceb11bc..e3d43acf7 100644
--- a/tests/test_util/test_io.py
+++ b/tests/test_util/test_io.py
@@ -295,22 +295,17 @@ def test_open_zip_multi(tmp_dir):
 
 
 @pytest.mark.parametrize(
-    "file_name,mode",
-    list(
-        product(
-            ("c_elegans_WS199_dna_shortened.fasta", "gff2_test.gff"),
-            ("r", "rb", "rt", None),
-        )
-    ),
+    "mode",
+    ("r", "rb", "rt", None),
 )
-def test_open_url(file_name, mode):
+def test_open_url(mode):
     # None value for open_url mode defaults to "rb"
-    local_root = pathlib.Path("data")
+    file_name = "gff2_test.gff"
     remote_root = (
         "https://raw.githubusercontent.com/cogent3/cogent3/develop/tests/data/{}"
     )
 
-    with open_(local_root / file_name, mode=mode) as infile:
+    with open_(DATADIR / file_name, mode=mode) as infile:
         local_data = infile.read()
 
     with open_url(remote_root.format(file_name), mode=mode) as infile:
@@ -323,6 +318,38 @@ def test_open_url(file_name, mode):
     assert remote_data.splitlines() == local_data.splitlines()
 
 
+def test_open_url_local():
+    """using file:///"""
+    # None value for open_url mode defaults to "rb"
+    file_name = "gff2_test.gff"
+    local_path = DATADIR / file_name
+    with open_(local_path) as infile:
+        local_data = infile.read()
+
+    # make absolute path
+    with open_url(local_path.absolute().as_uri()) as infile:
+        remote_data = infile.read()
+
+    assert remote_data.splitlines() == local_data.splitlines()
+
+
+def test_open_url_compressed():
+    """uses local compressed file to check"""
+    # None value for open_url mode defaults to "rb"
+    file_name = "formattest.fasta.gz"
+    remote_root = (
+        "https://raw.githubusercontent.com/cogent3/cogent3/develop/tests/data/{}"
+    )
+
+    with open_(DATADIR / file_name) as infile:
+        local_data = infile.read()
+
+    with open_url(remote_root.format(file_name)) as infile:
+        remote_data = infile.read()
+
+    assert remote_data.splitlines() == local_data.splitlines()
+
+
 def test_open_url_write_exceptions():
     """Test 'w' mode (should raise Exception)"""
     with pytest.raises(Exception):