from __future__ import annotations
import asyncio
import hashlib
import io
import json
import logging
from mgnipy._shared_helpers.httpx_helpers import init_httpx_client
logger = logging.getLogger(__name__)
from http.client import IncompleteRead
from pathlib import Path
from typing import Any, Generator, Literal, Optional
import webbrowser
import zlib
from bigtree import Tree
import httpx
import ijson
import pandas as pd
import polars as pl
from pydantic import HttpUrl
from skbio.io import read
from mgnipy._shared_helpers.writers import atomic_write_bytes, atomic_write_json
[docs]
class CheckpointMixin:
"""
Checkpoint manager for request-makers.
This mixin provides methods to write paginated results to disk as they are retrieved, and to load them back into memory on subsequent runs. It generates a unique cache key based on the query parameters and resource type, and organizes cached pages in a directory structure under a specified cache root.
The mixin assumes the host class has the following attributes:
- `self.params`: A dictionary of query parameters for the current request.
- `self.config`: An instance of `MGnipyConfig` containing configuration settings, including the cache directory.
- `self.resource`: A string representing the type of resource being queried (e.g., "samples", "runs").
- `self._results`: A dictionary mapping page numbers to lists of dictionaries containing the query results.
- `self.count`: An integer representing the total number of records for the current query.
- `self.num_requests`: An integer representing the total number of paginated requests needed for the current query.
which is the case for: QuerySets, MGnifier, proxies
Example
-------
>>> from mgnipy.V2.mgnifier.query_set import QuerySet
>>> from mgnipy.V2.mixins import CheckpointMixin
>>> # Creating a class that uses mixin
>>> class MyQuerySet(QuerySet, CheckpointMixin):
... # Initialize QuerySet (not mixin)
... def __init__(self, *args, **kwargs):
... super().__init__(*args, **kwargs)
>>> # init
>>> qs = MyQuerySet(resource="studies") # default config is w/ cache enabled
>>> # now can use checkpoint methods, e.g.
>>> qs.cache_key # print the unique cache key
'...'
>>> qs.load_cache()
[]
"""
@property
def cache_key(self) -> str:
"""
Generate deterministic hash from resource + params.
Returns
-------
str
A unique cache key for the current query parameters and resource.
For a query to the 'samples' resource with parameters {'biome_lineage': 'root:Environmental:Terrestrial'},
the cache key will be a SHA256 hash of the string representation of the resource and parameters,
ensuring that identical queries will have the same cache key and thus access the same cached results.
Example
-------
>>> from mgnipy.V2.mixins import CheckpointMixin
>>> from mgnipy import MGnipyConfig
>>> # Prepare parameters and config
>>> params = {'lineage': 'root:Environmental:Terrestrial'}
>>> resource = 'biome'
>>> config = MGnipyConfig(cache_dir="/path/to/cache")
>>> # Create CheckpointMixin instance and compute cache key
>>> cache_handler = CheckpointMixin()
>>> cache_handler.params = params
>>> cache_handler.resource = resource
>>> cache_handler.config = config
>>> cache_handler.cache_key
'1eb56ddf5a2e7d60d8155c8bbe01f032f959a2519d43e99f31f533abffa3166f'
"""
params: dict = self.params.copy()
serial: str = json.dumps(
{"resource": str(self.resource), "params": params},
sort_keys=True,
default=str,
)
cache_key: str = hashlib.sha256(serial.encode("utf-8")).hexdigest()
return cache_key
@property
def cache_path(self) -> Optional[Path]:
"""Directory for this query's cached pages."""
if self.config.cache_dir is None:
return None
return self.config.cache_dir / self.cache_key
@property
def manifest_path(self) -> Optional[Path]:
"""Path to mgnipy_manifest.json storing metadata."""
if self.cache_path is None:
return None
return self.cache_path / "mgnipy_manifest.json"
[docs]
def write_results(
self,
request_num: int,
items: Any,
include_manifest: bool = True,
) -> None:
"""Auto atomic write to disk."""
save_to = self.cache_path
if save_to is None:
logger.debug(f"Cache disabled: Skipping cache write for page {request_num}")
return
logger.debug(f"Creating cache dir if not exists: {save_to}")
save_to.mkdir(parents=True, exist_ok=True)
# prep filenames/paths
filepath = save_to / f"mgnipy_page_{request_num}.json"
logger.info(f"Writing page {request_num} to {filepath}")
# now actually writing the response results:
# Write bytes (binary downloads) using atomic_write_bytes, otherwise JSON
try:
if isinstance(items, (bytes, bytearray)):
bin_path = filepath.with_suffix(".bin")
atomic_write_bytes(bin_path, bytes(items))
else:
atomic_write_json(filepath, items)
except Exception:
logger.error(f"Failed to write cache file for page {request_num}")
if include_manifest:
manifest_path = self.manifest_path
logger.info(f"Writing manifest to {manifest_path}")
manifest = {
"resource": str(self.resource),
"params": self.params,
"count": self.count,
"total_pages": self.num_requests,
}
if manifest_path is not None:
atomic_write_json(manifest_path, manifest)
else:
logger.debug("Skipping manifest write for this page")
[docs]
async def awrite_results(self, request_num: int, items: Any) -> None:
"""Async wrapper for write_results."""
logger.debug(f"Asynchronously writing cached results for page {request_num}")
await asyncio.to_thread(self.write_results, request_num, items)
[docs]
def load_cache_results(self) -> list[int]:
"""
Load cached pages/request nums into results.
Loads cached pages from disk into the in-memory results dictionary (self._results), if available.
Returns
-------
list of int
A list of request numbers (page numbers) that were successfully loaded from the cache.
"""
# where to load from (or not)
load_from = self.cache_path
if load_from is None:
logger.debug(
f"Cache disabled: Skipping cache load for {str(self.resource)}."
)
return []
logger.info(f"Loading cached pages from {load_from}")
if not load_from.exists():
logger.info(f"No cache to load yet from {load_from}")
return []
# initialize stores
if self._results is None:
logger.debug("Initializing self._results as empty dict for cache load")
self._results = {}
pages_loaded = []
possible_pages = sorted(load_from.glob("mgnipy_page_*.*"))
logger.debug(f"Possible pages: {possible_pages}")
# load each page file
for cache_file in possible_pages:
if cache_file.suffix not in {".json", ".bin"}:
logger.info(f"Skipping {cache_file}")
continue
# Extract page number from filename
try:
request_num = int(cache_file.stem.split("_")[-1])
except Exception as e:
logger.error(f"Failed to extract page number from {cache_file}: {e}")
continue
if request_num in self._results:
logger.debug(
f"Page {request_num} already loaded in memory; skipping {cache_file}"
)
continue
logger.info(f"Loading {cache_file}")
try:
if cache_file.suffix == ".bin":
with cache_file.open("rb") as fh:
data = fh.read()
else:
with cache_file.open("r", encoding="utf-8") as fh:
data = json.load(fh)
# logger.debug(f"Loaded JSON data for page {request_num}: {data}")
# load page to results
self._results[request_num] = data
logger.debug(f"Loaded page {request_num} to self._results")
# tracking
pages_loaded.append(request_num)
logger.debug(f"Pages loaded so far: {pages_loaded}")
except Exception as e:
logger.warning(f"Failed to load cache file: {cache_file}. Error: {e}")
return pages_loaded
[docs]
def load_cache_manifest(self) -> dict:
"""
Load the cache manifest file if present, and update total records and total requests.
Returns
-------
dict
The contents of the manifest file, or an empty dictionary if the manifest is not found or fails to load.
"""
# Load manifest if present
mpath = self.manifest_path
if mpath is None:
return {}
if mpath.exists():
logger.info(f"Loading cache manifest from {mpath}")
try:
with mpath.open("r", encoding="utf-8") as fh:
manifest = json.load(fh)
self.count = manifest.get("count")
self.num_requests = manifest.get("total_pages")
except Exception:
logger.warning(f"Failed to load manifest file: {mpath}")
manifest = {}
else:
manifest = {}
return manifest
[docs]
def load_cache(self) -> list[int]:
"""
Pick up where you left off. Loads cached results and manifest into memory.
Returns
-------
list of int
A list of request numbers (page numbers) that were successfully loaded from the cache.
"""
load_from = self.cache_path
if load_from is None:
logger.debug("Skipping cache load because cache is disabled")
return []
logger.info(f"Loading cache from {self.cache_path}")
pages_loaded = self.load_cache_results()
mani = self.load_cache_manifest()
logger.info(f"Loaded {len(pages_loaded)} cached pages and manifest: {mani}")
return pages_loaded
[docs]
async def aload_cache(self) -> list[int]:
"""Async wrapper for load_cache."""
return await asyncio.to_thread(self.load_cache)
[docs]
def clear_cache(self) -> None:
"""Remove all cached pages for this set of queries."""
load_from = self.cache_path
if load_from is None:
logger.info("Cache is disabled; skipping cache clearing")
return
if load_from.exists():
logger.info(f"Clearing cache directory {load_from}")
for cache_file in load_from.iterdir():
# extra check just in case
if cache_file.name == "mgnipy_manifest.json" or (
cache_file.name.startswith("mgnipy_page_")
and cache_file.suffix in {".json", ".bin"}
):
try:
logger.debug(f"Deleting cache file {cache_file}")
cache_file.unlink()
except Exception:
logger.warning(f"Failed to delete cache file: {cache_file}")
try:
load_from.rmdir()
except Exception:
logger.warning(f"Failed to delete cache directory: {load_from}")
# reset loaded cache state
self._pages_from_cache = []
self._cached_manifest = {}
[docs]
def try_load_cache(self) -> None:
"""
Attempt to load cached results and manifest into memory if not already loaded.
This method checks if the cache has already been loaded to avoid redundant operations.
If the cache has not been loaded, it will attempt to load it and set the `_cache_loaded` attribute accordingly.
Notes
-----
- This method is intended to be called internally before accessing cached results.
- If cache_dir is None then _cache_loaded will be True after initial attempt.
- If an error occurs during cache loading, it will be logged, and `_cache_loaded` will be set to False.
- Dependent on `.mixins.CheckpointMixin`
"""
if getattr(self, "_cache_loaded", False):
logger.debug(
f"Cache already loaded; skipping load attempt. "
f"Pages loaded: {self._pages_from_cache}"
f"Manifest: {self._cached_manifest}"
)
return
try:
# load results
self._pages_from_cache = self.load_cache_results()
# load manifest
self._cached_manifest = self.load_cache_manifest()
# update counts from manifest
self.count = self._cached_manifest.get("count", None)
self.num_requests = self._cached_manifest.get("total_pages", None)
# set flag
self._cache_loaded = True
except Exception as e:
logger.error(f"Error occurred while loading cache: {e}")
self._cache_loaded = (
False # Q: Or should this be set to True to avoid repeated attempts?
)
[docs]
class StreamMixin:
"""
Mixin providing streaming helpers for downloads.
# TODO remove below dependencies on mgnifier
This mixin assumes the host class provides the following helpers/properties:
`.exec.httpx_client` and `.exec.httpx_aclient` attributes
- `_get_type_by_alias(alias)` to resolve file types
- `downloads_df` when needed for examples/tests
The implementation mirrors the streaming helpers previously defined
on :class:`MGazine` so they can be reused by other classes.
"""
def _handle_incomplete_read(self, url: str):
# TODO
logger.warning(
f"You can also download the file {url} using the 'download' method instead of streaming and then read it into memory from disk, which may be more reliable for unstable connections."
)
raise IncompleteRead(
f"Incomplete read error encountered when streaming {url}. This may be due to a network issue or server timeout. Consider retrying the request or checking your connection."
) from None
[docs]
def stream_pandas(
self,
url: str,
sep: str = "\t",
chunksize: Optional[int] = None,
max_skip: int = 5,
low_memory: bool = False,
**pd_kwargs,
) -> pd.DataFrame | pd.io.parsers.readers.TextFileReader:
"""
Read a TSV from a URL or local file with resilient header handling.
The helper will retry with increasing ``skiprows`` when ``pandas``
raises a ``ParserError`` (useful for files with extra header lines).
When ``chunksize`` is provided an iterator is returned.
Parameters
----------
url : str
The URL or local file path to read the TSV from.
sep : str
The delimiter to use (default is tab).
chunksize : int or None
If an integer is provided, returns an iterator that yields DataFrames
of that many rows. If None, returns a single DataFrame.
max_skip : int
The maximum number of lines to skip when trying to parse the TSV.
**pd_kwargs
Additional keyword arguments passed to ``pd.read_csv``.
Returns
-------
pd.DataFrame or TextFileReader
A DataFrame containing the TSV data, or an iterator yielding DataFrames
if ``chunksize`` is specified.
Raises
------
ValueError
If ``chunksize`` is not a positive integer or None.
RuntimeError
If the TSV cannot be parsed after skipping up to ``max_skip`` lines.
Pandas ParserError
If the TSV cannot be parsed due to a format error (after retries).
"""
for skip in range(max_skip + 1):
try:
return pd.read_csv(
url,
sep=sep,
chunksize=chunksize,
skiprows=skip if skip > 0 else None,
low_memory=low_memory,
**pd_kwargs,
)
except pd.errors.ParserError:
continue # Try next skiprows value
except IncompleteRead:
self._handle_incomplete_read(url)
except Exception as err:
raise err
raise pd.errors.ParserError(
f"Failed to parse {url} after skipping up to {max_skip} rows."
)
[docs]
def stream_polars(
self,
url: str,
sep: str = "\t",
chunksize: Optional[int] = None,
max_skip: int = 5,
low_memory: bool = False,
**pl_kwargs,
) -> pl.DataFrame | pl.LazyFrame:
"""
Read a TSV from a URL or local file into a Polars DataFrame with resilient header handling.
The helper will retry with increasing ``skip_rows`` when Polars raises an error
(useful for files with extra header lines). When ``chunksize`` is provided an iterator is returned.
Parameters
----------
url : str
The URL or local file path to read the TSV from.
sep : str
The delimiter to use (default is tab).
chunksize : int or None
If an integer is provided, returns an iterator that yields DataFrames
of that many rows. If None, returns a single DataFrame.
max_skip : int
The maximum number of lines to skip when trying to parse the TSV.
**pl_kwargs
Additional keyword arguments passed to ``pl.read_csv``.
Returns
-------
pl.DataFrame or Iterator[pl.DataFrame]
A Polars DataFrame containing the TSV data, or an iterator yielding DataFrames
if ``chunksize`` is specified.
Raises
------
ValueError
If ``chunksize`` is not a positive integer or None.
RuntimeError
If the TSV cannot be parsed after skipping up to ``max_skip`` lines.
Polars Error
If the TSV cannot be parsed due to a format error (after retries).
"""
if chunksize is None:
the_chosen = pl.read_csv
else:
the_chosen = pl.scan_csv
for skip in range(max_skip + 1):
try:
return the_chosen(
url,
separator=sep,
skip_rows_after_header=skip,
truncate_ragged_lines=True,
infer_schema_length=10000,
low_memory=low_memory,
**pl_kwargs,
)
except pl.exceptions.PolarsError:
continue # Try next skip_rows value
except IncompleteRead:
self._handle_incomplete_read(url)
except Exception as err:
raise err
raise pl.exceptions.PolarsError(
f"Failed to parse {url} after skipping up to {max_skip} rows."
)
[docs]
def stream_html(self, url: str, **web_kwargs) -> bool:
"""
Open an HTML URL in the default web browser.
Parameters
----------
url : str
The URL to open in the web browser.
**web_kwargs
Additional keyword arguments passed to `webbrowser.open()`, such as `new` and `autoraise`.
Returns
-------
bool
True if the URL was opened successfully, False otherwise.
"""
return webbrowser.open(url, **web_kwargs)
[docs]
def stream_txt(
self,
url: str,
chunksize: Optional[int] = None,
**httpx_kwargs,
) -> str | Generator:
"""
Stream a plain-text resource.
When ``chunksize`` is ``None`` the full text is returned as a string.
When ``chunksize`` is an integer the function yields lists of lines.
Parameters
----------
url : str
The URL to stream the text from.
chunksize : int or None
If an integer is provided, yields lists of lines of that size. If None, yields the entire text as a single string.
httpx_client : httpx.Client, optional
An optional httpx.Client to use for the request. If None, a new client will be created for the request.
**httpx_kwargs
Additional keyword arguments passed to the httpx.Client.request() method
Returns
-------
str or Generator
The full text as a string if chunksize is None, or a generator yielding lists of lines if chunksize is an integer.
"""
if chunksize is None:
# load as whole
with self.httpx_client.get(url, **httpx_kwargs) as response:
response.raise_for_status()
return response.text
elif isinstance(chunksize, int) and chunksize > 0:
# load in chunks
with self.httpx_client.stream("GET", url, **httpx_kwargs) as response:
response.raise_for_status()
chunk = []
for line in response.iter_text():
chunk.append(line)
if len(chunk) == chunksize:
yield chunk
chunk = []
if chunk:
yield chunk
else:
raise ValueError("`chunksize` must be a positive integer or None.")
[docs]
def stream_fasta(self, url: str, **skbio_kwargs) -> Generator:
"""
Stream a FASTA file from a URL using scikit-bio's read function. Refer there for more info.
Parameters
----------
url : str
The URL to the FASTA file to stream.
**skbio_kwargs
Additional keyword arguments passed to `skbio.io.read()`, such as `into` and `verify`.
Returns
-------
Generator
A generator yielding scikit-bio Sequence objects parsed from the FASTA file.
"""
return read(url, format="fasta", **skbio_kwargs)
[docs]
def stream_gff(self, url: str, **skbio_kwargs) -> Generator:
"""
Stream a GFF file from a URL using scikit-bio's read function. Refer there for more info.
Parameters
----------
url : str
The URL to the GFF file to stream.
**skbio_kwargs
Additional keyword arguments passed to `skbio.io.read()`, such as `into` and `verify`.
Returns
-------
Generator
A generator yielding scikit-bio Sequence objects parsed from the GFF file.
"""
return read(url, format="gff3", **skbio_kwargs)
[docs]
def stream_biom(self, url: str, **skbio_kwargs) -> Generator:
"""
Stream a biom file from a URL using scikit-bio's read function. Refer there for more info.
Parameters
----------
url : str
The URL to the biom file to stream.
**skbio_kwargs
Additional keyword arguments passed to `skbio.io.read()`, such as `into` and `verify`.
Returns
-------
Generator
A generator yielding scikit-bio Sequence objects parsed from the biom file.
"""
return read(url, format="biom", **skbio_kwargs)
[docs]
def stream_gzipped(
self,
url: str,
chunksize: Optional[int] = None,
decode: bool = False,
encoding: str = "utf-8",
errors: str = "replace",
**httpx_kwargs,
) -> bytes | str | io.BufferedReader | io.TextIOWrapper:
"""Stream a gzipped HTTP resource and present a file-like interface.
When ``chunksize`` is None the entire compressed payload is fetched
and decompressed into memory. When ``chunksize`` is provided a
streaming file-like object is returned.
"""
logger.debug(
"stream_gzipped called url=%s chunksize=%s decode=%s",
url,
chunksize,
decode,
)
if chunksize is None:
logger.debug("Using full-download mode (chunksize=None)")
r = self.httpx_client.get(url, timeout=None, **httpx_kwargs)
r.raise_for_status()
decompressor = zlib.decompressobj(16 + zlib.MAX_WBITS)
data = decompressor.decompress(r.content) + decompressor.flush()
logger.debug(
"Full-download mode complete: compressed=%d decompressed=%d",
len(r.content),
len(data),
)
return data.decode(encoding, errors=errors) if decode else data
if not isinstance(chunksize, int) or chunksize <= 0:
raise ValueError("`chunksize` must be a positive integer or None.")
class _HTTPGzipRaw(io.RawIOBase):
def __init__(self):
self._cm = self.httpx_client.stream(
"GET", url, timeout=None, **httpx_kwargs
)
self._resp = self._cm.__enter__()
self._resp.raise_for_status()
self._iter = self._resp.iter_raw(chunk_size=chunksize)
self._decomp = zlib.decompressobj(16 + zlib.MAX_WBITS)
self._buf = bytearray()
self._eof = False
self._flushed = False
logger.debug("Streaming HTTP/gzip reader initialized")
def readable(self) -> bool:
return True
def _fill(self, need: int) -> None:
while len(self._buf) < need and not self._eof:
try:
chunk = next(self._iter)
except StopIteration:
if not self._flushed:
tail = self._decomp.flush()
if tail:
self._buf.extend(tail)
self._flushed = True
self._eof = True
logger.debug("Reached end of HTTP stream")
break
if chunk:
out = self._decomp.decompress(chunk)
if out:
self._buf.extend(out)
def readinto(self, b) -> int:
if self.closed:
return 0
mv = memoryview(b)
self._fill(len(mv))
n = min(len(mv), len(self._buf))
if n <= 0:
return 0
mv[:n] = self._buf[:n]
del self._buf[:n]
return n
def close(self) -> None:
if not self.closed:
try:
self._cm.__exit__(None, None, None)
finally:
super().close()
logger.debug("Streaming HTTP/gzip reader closed")
raw = _HTTPGzipRaw()
buffered = io.BufferedReader(raw, buffer_size=chunksize)
if decode:
return io.TextIOWrapper(buffered, encoding=encoding, errors=errors)
return buffered
[docs]
def stream_jsonl(
self,
url: str,
orient: Optional[
Literal["records", "split", "index", "columns", "values", "table"]
] = None,
chunksize: Optional[int] = None,
df_engine: Optional[Literal["pandas", "polars"]] = "pandas",
**df_kwargs,
) -> dict:
if df_engine == "pandas":
return pd.read_json(
url, orient=orient, lines=True, chunksize=chunksize, **df_kwargs
)
elif df_engine == "polars":
if chunksize is None:
return pl.read_ndjson(url, infer_schema_length=10000, **df_kwargs)
else:
return pl.scan_ndjson(url, infer_schema_length=10000, **df_kwargs)
[docs]
def stream_json(
self,
url: str,
chunksize: Optional[int] = None,
**httpx_kwargs,
) -> dict | Generator:
if chunksize is None and not (url.endswith(".gz") or url.endswith(".gzip")):
with self.httpx_client.get(url, **httpx_kwargs) as response:
response.raise_for_status()
return response.json()
elif chunksize is not None and not (
url.endswith(".gz") or url.endswith(".gzip")
):
with self.httpx_client.stream("GET", url, **httpx_kwargs) as response:
response.raise_for_status()
for entry in ijson.kvitems(response.iter_text(), ""):
yield entry
elif url.endswith(".gz") or url.endswith(".gzip"):
with self.stream_gzipped(
url,
chunksize=chunksize,
decode=True,
**httpx_kwargs,
) as gzipped_stream:
for entry in ijson.kvitems(gzipped_stream, ""):
yield entry
else:
raise ValueError(f"Unsupported file type for URL: {url}")
def _fix_inconsistent_cols(
self, fields: list[str], pad_to: int = 15
) -> list[str] | None:
"""Pad or truncate list of fields to ``pad_to`` length.
Parameters
----------
fields : list of str
List of column names to adjust.
pad_to : int, optional
Desired length of the returned list. Defaults to 15.
Returns
-------
list of str or None
The adjusted list of fields or ``None`` when ``pad_to`` is 0.
"""
if len(fields) < pad_to:
return fields + [""] * (pad_to - len(fields))
if len(fields) > pad_to:
return fields[:pad_to]
return fields
[docs]
def stream_tree(self, url: str, **skbio_kwargs) -> Generator:
return read(url, format="newick", **skbio_kwargs)
def _get_streamer(
self,
alias: Optional[str] = None,
url: Optional[HttpUrl] = None,
chunksize: int = 1000,
max_skip: int = 5,
df_engine: Optional[Literal["pandas", "polars"]] = "pandas",
low_memory: bool = False,
**kwargs,
):
_alias, _url = self._prioritize_alias(alias, url, required=True)
file_type = self._get_type_by_alias(_alias)
if df_engine == "polars" and file_type == "tsv":
return self.stream_polars(
_url,
sep="\t",
chunksize=chunksize,
max_skip=max_skip,
low_memory=low_memory,
**kwargs,
)
if df_engine == "polars" and file_type == "csv":
return self.stream_polars(
_url,
sep=",",
chunksize=chunksize,
max_skip=max_skip,
low_memory=low_memory,
**kwargs,
)
if file_type == "tsv":
if _url.endswith(".gz") or _url.endswith(".gzip"):
logger.debug(f"tsv file type ends with .gz: {_url}")
try:
return self.stream_pandas(
_url,
chunksize=chunksize,
max_skip=max_skip,
compression="gzip",
**kwargs,
)
except pd.errors.ParserError as e:
logger.error(f"ParserError: {e}")
return self.stream_pandas(
_url,
chunksize=chunksize,
max_skip=max_skip,
compression="gzip",
engine="python",
on_bad_lines=self._fix_inconsistent_cols,
**kwargs,
)
elif _url.endswith(".txt") or _url.endswith(".tsv"):
return self.stream_pandas(
_url,
chunksize=chunksize,
max_skip=max_skip,
low_memory=low_memory,
**kwargs,
)
if file_type == "csv":
return self.stream_pandas(
_url,
sep=",",
chunksize=chunksize,
max_skip=max_skip,
low_memory=low_memory,
**kwargs,
)
if file_type == "html":
return lambda: self.stream_html(_url, **kwargs)
if file_type == "txt":
return self.stream_txt(_url, chunksize=chunksize, **kwargs)
if file_type == "gff":
return self.stream_gff(_url, **kwargs)
if file_type == "biom":
return self.stream_biom(_url, **kwargs)
if file_type == "fasta":
return self.stream_fasta(_url, **kwargs)
if file_type == "tree":
return self.stream_tree(_url, **kwargs)
if file_type == "json":
return self.stream_jsonl(
_url,
orient="records",
chunksize=chunksize,
df_engine=df_engine,
low_memory=low_memory,
**kwargs,
)
if file_type == "other" and ".json" in _url:
if _url.endswith("json.gz") or _url.endswith("json.gzip"):
return self.stream_json(_url, chunksize=chunksize, **kwargs)
logger.info(
f"{_alias} is only available for download (e.g., `.download({_alias}))`"
)
logger.debug(
f"Alias: {_alias}\nURL: {_url}\nFile type: {file_type}. Only '.json' files can be streamed under 'other' type, otherwise this download is only available for download."
)
else:
raise ValueError(f"Unsupported file type for streaming: {file_type}")
[docs]
def stream(
self,
*,
alias: Optional[str] = None,
url: Optional[HttpUrl] = None,
chunksize: Optional[int] = None,
max_skip: int = 5,
**kwargs,
) -> Any:
"""
Streams a single download based on its alias or url.
If ``chunksize`` is specified then iterators of dataframes or strings
will be returned; otherwise the full data will be returned as a single
object.
Supported formats and their handlers
------------------------------------
- tsv: handled by :meth:`stream_pandas` (pandas) or :meth:`stream_polars` (polars).
Gzipped TSVs are supported via the gzip/compression options.
- csv: handled by :meth:`stream_pandas` / :meth:`stream_polars` (sep=",").
- txt: handled by :meth:`stream_txt` (returns full text or yields line chunks).
- html: handled by :meth:`stream_html` (opens URL in browser).
- fasta: handled by :meth:`stream_fasta` (scikit-bio generator).
- gff: handled by :meth:`stream_gff` (scikit-bio generator).
- biom: handled by :meth:`stream_biom` (scikit-bio generator).
- gzipped HTTP resources: use :meth:`stream_gzipped` for a file-like object,
or :meth:`stream_json` for gzipped JSON content.
- jsonl / ndjson: handled by :meth:`stream_jsonl` (pandas or polars modes).
- json: handled by :meth:`stream_json` (returns full JSON or streams via ijson).
- tree/newick: handled by :meth:`stream_tree` (scikit-bio newick reader).
- other: if the URL ends with ``.json`` it's streamed via :meth:`stream_json`;
otherwise use the download helper for unsupported binary formats.
Parameters
----------
alias : Optional[str]
The alias of the download to stream.
url : Optional[HttpUrl]
The url of the download to stream.
chunksize : Optional[int]
The size of the chunks to read from the stream.
max_skip : int, optional
The maximum number of rows to skip before raising an error. Default is 5.
**kwargs
Additional keyword arguments to pass to the streamer function.
Returns
-------
Any
The streamer result for the resolved alias or url.
"""
# resolve a single alias/url target
_alias, _url = self._prioritize_alias(alias, url, required=True)
# return a single streamer result, not a dict of all streams
logger.info("Setting up stream for alias=%s url=%s", _alias, _url)
try:
return self._get_streamer(
alias=_alias,
url=_url,
chunksize=chunksize,
max_skip=max_skip,
**kwargs,
)
except httpx.HTTPError as err:
logger.error("HTTP error for alias=%s url=%s: %s", _alias, _url, err)
raise
[docs]
class BiomesTreeMixin:
@property
def lineages(self) -> list[str]:
mgnify_metadata = self.search_results
return mgnify_metadata.ids
@property
def tree(self) -> Tree:
"""
Convert the biomes metadata to a tree structure for visualization or analysis.
Returns
-------
Tree
A tree representation of the biomes and their relationships.
"""
logger.debug("Building tree from %s lineages", len(self.lineages))
# TODO generate nodes first
return Tree.from_list(self.lineages, sep=":")
[docs]
def show_tree(
self,
method: Literal[
"compact",
"show",
"print",
"horizontal",
"hshow",
"h",
"hprint",
"vertical",
"vshow",
"v",
"vprint",
] = "compact",
):
logger.info("Showing tree using method %s", method)
if method in ["compact", "show", "print"]:
# TODO print_tree(self._tree)
self.tree.show()
elif method in ["horizontal", "hshow", "h", "hprint"]:
self.tree.hshow()
elif method in ["vertical", "vshow", "v", "vprint"]:
self.tree.vshow()
else:
raise ValueError(
f"Invalid method: {method}. "
"Supported methods: 'compact', 'show', 'print', "
"'horizontal', 'hshow', 'h', 'hprint', "
"'vertical', 'vshow', 'v', 'vprint'."
)
@property
def results(self) -> dict:
"""Get results and auto-normalize lineage field."""
parent_results = super().results
# Always normalize if results exist
if parent_results:
logger.debug("Normalizing lineage fields in results")
self._normalise_lineage()
return parent_results
def _normalise_lineage(self):
"""
Rename field "lineage" to "biome_lineage" for consistency with other resources.
"""
if self._results:
logger.debug("Renaming lineage fields to biome_lineage")
for page_data in self._results.values():
if isinstance(page_data, list):
for record in page_data:
if isinstance(record, dict) and "lineage" in record:
record["biome_lineage"] = record.pop("lineage")
[docs]
class ClientManagerMixin:
"""
Context manager mixin for managing the lifecycle of an HTTP client (synchronous or asynchronous) within a class. This mixin provides methods to enter and exit the context, ensuring that the client is properly initialized and closed, and handles cases where the client may have been closed or is not available.
Requirements
------------
- `self.client` must be an instance of a class that implements the context manager protocol (i.e., has `__enter__` and `__exit__` methods for synchronous clients, and `__aenter__` and `__aexit__` methods for asynchronous clients).
- `self.config` must be a configuration object or dictionary that can be used to initialize a new HTTP client instance if the existing client is closed or unavailable.
"""
def __enter__(self):
try:
self.client.__enter__()
except RuntimeError as e:
logger.warning(f"Opening a new client instance due to: {e}")
self.client = init_httpx_client(self.config)
self.client.__enter__()
return self
def __exit__(self, *args, **kwargs):
if not getattr(self, "_owns_client", True):
logger.debug("Client is not owned by this instance; skipping __exit__().")
return False
self.client.__exit__(*args, **kwargs)
self.close()
logger.debug(f"Closing client: {self.client._client}")
return False
async def __aenter__(self):
try:
await self.client.__aenter__()
except RuntimeError as e:
logger.warning(f"Opening a new async client instance due to: {e}")
self.client = init_httpx_client(self.config)
await self.client.__aenter__()
return self
async def __aexit__(self, *args, **kwargs):
if not getattr(self, "_owns_client", True):
logger.debug("Client is not owned by this instance; skipping __aexit__().")
return False
await self.client.__aexit__(*args, **kwargs)
await self.aclose()
logger.debug(f"Closing async client: {self.client._async_client}")
return False
[docs]
def close(self):
if self.client is not None and self.client._client is not None:
if not self.client._client.is_closed:
self.client._client.close()
self.client._client = None
[docs]
async def aclose(self):
if self.client is not None and self.client._async_client is not None:
if not self.client._async_client.is_closed:
await self.client._async_client.aclose()
self.client._async_client = None
[docs]
def status(self) -> None:
"""
Print the status of the MGnipy client, including the type of client and whether the synchronous and asynchronous httpx client sessions are open.
"""
print(
f"Client or AuthenticatedClient type: {type(self.client).__name__}\n"
f"HTTP client open: {self.client._client is not None and not self.client._client.is_closed}\n"
f"Async client open: {self.client._async_client is not None and not self.client._async_client.is_closed}\n"
)
[docs]
def renew_client(self):
"""
Init a new client instance and replace the existing one. This is useful if the current client has been closed or is no longer valid, allowing for a fresh start with a new HTTP client session.
"""
logger.info(f"Renewing HTTP client instance. Old: {self.client}")
self.client = init_httpx_client(self.config)
logger.info(f"HTTP client refreshed. Now: {self.client}")
@property
def httpx_client(self) -> httpx.Client:
"""
Get the synchronous httpx client instance from the AuthenticatedClient.
Returns
-------
httpx.Client
The synchronous httpx client instance.
"""
if self.client is None:
logger.info("Client is None; initializing a new client instance.")
self.renew_client()
elif self.client._client is None:
logger.info(
"Synchronous httpx Client._client is None. Getting httpx client"
)
elif self.client._client.is_closed:
logger.info("Synchronous httpx Client._client is closed; renewing client.")
self.renew_client()
self.client.get_httpx_client()
return self.client._client
@property
def async_httpx_client(self) -> httpx.AsyncClient:
"""
Get the asynchronous httpx client instance from the AuthenticatedClient.
Returns
-------
httpx.AsyncClient
The asynchronous httpx client instance.
"""
if self.client is None:
logger.info("Client is None; initializing a new client instance.")
self.renew_client()
elif self.client._async_client is None:
logger.info(
"Asynchronous httpx Client._async_client is None. Getting async httpx client"
)
elif self.client._async_client.is_closed:
logger.info(
"Asynchronous httpx Client._async_client is closed; renewing client."
)
self.renew_client()
self.client.get_httpx_async_client()
return self.client._async_client
