"""Serialize and deserialize :class:`KinaseInfo` objects to/from json/yaml/toml and tar archives.
Provides :func:`serialize_kinase_dict` and :func:`deserialize_kinase_dict` for
round-tripping per-kinase files, plus helpers for reading packaged ``.tar.gz`` data
in memory. TOML serialization is skipped on Windows.
"""
import glob
import json
import logging
import os
import shutil
import tarfile
from importlib import resources
from io import BytesIO
from typing import Any, Optional
import git
import toml
import yaml
from mkt.schema import kinase_schema
from mkt.schema.config import get_output_dir
from mkt.schema.utils import TQDM_BAR_FORMAT
from pydantic import BaseModel
from tqdm import tqdm
logger = logging.getLogger(__name__)
_deserialization_cache = {}
[docs]
def get_repo_root():
"""Get the root of the git repository.
Returns
-------
str
Path to the root of the git repository; if not found, return current directory
"""
try:
repo = git.Repo(".", search_parent_directories=True)
return repo.working_tree_dir
except git.InvalidGitRepositoryError:
logger.info("Not a git repository; using current directory as root...")
return "."
DICT_FUNCS = {
"json": {
"serialize": json.dumps,
"kwargs_serialize": {"default": list, "indent": 4},
"deserialize_file": json.load,
"deserialize_str": json.loads,
"kwargs_deserialize": {},
},
"yaml": {
"serialize": yaml.safe_dump,
"kwargs_serialize": {"sort_keys": False},
"deserialize_file": yaml.safe_load,
"deserialize_str": yaml.safe_load,
"kwargs_deserialize": {},
},
"toml": {
"serialize": toml.dumps,
"kwargs_serialize": {},
"deserialize_file": toml.load,
"deserialize_str": toml.loads,
"kwargs_deserialize": {},
},
}
"""dict[str, dict[str, Callable]]: Dictionary of serialization and deserialization functions supported."""
[docs]
def return_filenotfound_error_if_empty_or_missing(
str_path_in: str,
) -> FileNotFoundError | None:
"""Return FileNotFoundError for the given path.
Parameters
----------
str_path_in : str
Path that was not found.
Returns
-------
FileNotFoundError | None
FileNotFoundError for the given path. If the path is not empty, return None.
"""
if os.path.exists(str_path_in) and str_path_in.endswith(".tar.gz"):
logger.info(
f"File {str_path_in} exists as a tar.gz directory. Will extract in memory..."
)
return None
elif not os.path.exists(str_path_in) or len(os.listdir(str_path_in)) == 0:
return FileNotFoundError
elif os.path.exists(str_path_in) and str_path_in.endswith(".tar.gz"):
logger.info(
f"File {str_path_in} exists as a tar.gz directory. Will extract in memory..."
)
return None
else:
return None
[docs]
def untar_if_neeeded(str_filename: str) -> str:
"""Unzip the file if it is a zip file.
Parameters
----------
str_filename : str
Path to the file.
Returns
-------
str
Path to the unzipped file or original file if not .tar.gz.
"""
if str_filename.endswith(".tar.gz"):
str_path_extract = os.path.dirname(str_filename)
extract_tarfiles(str_filename, str_path_extract)
str_filename = str_filename.replace(".tar.gz", "")
return str_filename
[docs]
def untar_files_in_memory(
str_path: str,
bool_extract: bool = True,
list_ids: list[str] | None = None,
) -> dict[str, str]:
"""Untar files exclusively in memory.
Parameters
----------
str_path : str
Path to the tar.gz file.
bool_extract : bool, optional
If True, extract the files to memory, by default True.
list_ids : list[str] | None, optional
List of IDs to filter the files if reading from memory, by default None.
If None, all files will be extracted.
Returns
-------
dict[str, str]
Dictionary of file names and their contents as strings.
"""
with open(str_path, "rb") as f:
tar_data = f.read()
list_entries, dict_bytes = [], {}
with BytesIO(tar_data) as tar_buffer, tarfile.open(
fileobj=tar_buffer, mode="r"
) as tar:
for member in tar.getmembers():
filename = os.path.basename(member.name)
# make sure entry is file
cond1 = member.isfile()
# ignore MacOS AppleDouble files
cond2 = "._" not in filename
# use list_ids, if provided
cond3 = list_ids is None or filename.split(".")[0] in list_ids
if cond1 and cond2 and cond3:
list_entries.append(filename.split(".")[0])
if bool_extract:
with tar.extractfile(member) as f:
dict_bytes[member.name] = f.read()
if bool_extract:
# decode bytes to string
dict_bytes = {k: v.decode("utf-8") for k, v in dict_bytes.items()}
return list_entries, dict_bytes
[docs]
def return_str_path_from_pkg_data(
str_path: str | None = None,
pkg_name: str | None = None,
pkg_resource: str | None = None,
) -> str:
"""Return the path to the package data directory or of a user-provided directory.
Parameters
----------
str_path : str | None, optional
Path to the KinaseInfo directory, by default None.
pkg_name : str | None, optional
Package name, by default None and will use mkt.schema.
pkg_resource : str | None, optional
Package resource, by default None and will use KinaseInfo.
Returns
-------
str
Path to the package data or user-provided directory.
"""
if pkg_name is None:
pkg_name = "mkt.schema"
if pkg_resource is None:
pkg_resource = "KinaseInfo.tar.gz"
if str_path is None:
try:
str_path = os.path.join(resources.files(pkg_name), pkg_resource)
# str_path = untar_if_neeeded(str_path)
return_filenotfound_error_if_empty_or_missing(str_path)
except Exception as e:
logger.error(
f"Could not find {pkg_resource} directory within {pkg_name}: {e}"
f"\nPlease provide a path to the {pkg_resource} directory."
)
else:
if not os.path.exists(str_path):
os.makedirs(str_path)
return str_path
[docs]
def clean_files_and_delete_directory(list_files: list[str]) -> None:
"""Remove unzipped files.
Parameters
----------
list_files : list[str]
List of files to remove.
Returns
-------
None
None
"""
try:
paths_remove = {os.path.dirname(i) for i in list_files}
[shutil.rmtree(i) for i in paths_remove if os.path.isdir(i)]
logger.info(f"Removed unzipped files: {[i for i in paths_remove]}.")
except Exception as e:
logger.error(f"Exception {e}")
logger.info(f"Could not remove unzipped files: {list_files}.")
# adapted from https://axeldonath.com/scipy-2023-pydantic-tutorial/notebooks-rendered/4-serialisation-and-deserialisation.html
[docs]
def serialize_kinase_dict(
kinase_dict: dict[str, BaseModel],
suffix: str = "json",
serialization_kwargs: Optional[dict[str, Any]] = None,
str_path: str | None = None,
) -> None:
"""Serialize KinaseInfo object to files.
Parameters
----------
kinase_dict : dict[str, BaseModel]
Dictionary of KinaseInfo objects.
suffix : str
Serialization types supported: json, yaml, toml.
serialization_kwargs : dict[str, Any], optional
Additional keyword arguments for serialization function, by default None;
(e.g., {"indent": 2} for json.dumps, {"sort_keys": False} for yaml.safe_dump).
str_path: str | None = None
Path to save the serialized file, by default None will use package data or Github repo data.
"""
if suffix not in DICT_FUNCS:
logger.error(
f"Serialization type ({suffix}) not supported; must be json, yaml, or toml."
)
return None
if os.name == "nt" and suffix == "toml":
logger.info("TOML serialization is not supported on Windows.")
return None
if serialization_kwargs is None:
serialization_kwargs = DICT_FUNCS[suffix]["kwargs_serialize"]
str_path = return_str_path_from_pkg_data(str_path)
for key, val in tqdm(
kinase_dict.items(),
desc="Serializing KinaseInfo objects...",
bar_format=TQDM_BAR_FORMAT,
):
with open(f"{str_path}/{key}.{suffix}", "w") as outfile:
val_serialized = DICT_FUNCS[suffix]["serialize"](
val.model_dump(),
**serialization_kwargs,
)
outfile.write(val_serialized)
[docs]
def deserialize_kinase_dict(
suffix: str = "json",
deserialization_kwargs: Optional[dict[str, Any]] = None,
str_path: str | None = None,
bool_remove: bool = True,
list_ids: list[str] | None = None,
str_name: str | None = None,
bool_verbose: bool = True,
) -> dict[str, BaseModel]:
"""Deserialize KinaseInfo object from files.
Parameters
----------
suffix : str
Deserialization types supported: json, yaml, toml.
deserialization_kwargs : dict[str, Any], optional
Additional keyword arguments for deserialization function, by default None.
str_path : str | None, optional
Path from which to load files, by default None.
bool_remove : bool, optional
If True, remove the files after deserialization, by default True.
list_ids : list[str] | None, optional
List of IDs to filter the files if reading from memory, by default None.
str_name : str | None, optional
Name of a variable in the global scope that contains the kinase dictionary to prevent reloading, by default None.
Returns
-------
dict[str, KinaseInfo]
Dictionary of KinaseInfo objects.
"""
if str_name is not None and str_name in _deserialization_cache:
if bool_verbose:
logger.info(f"Loading KinaseInfo object from variable {str_name}...")
return _deserialization_cache[str_name]
# callers that only need the modules importable (e.g. a Sphinx docs build or
# a CLI printing usage text) can set MKT_SKIP_KINASE_DICT to skip the
# expensive deserialization; cache the empty result so repeat calls are cheap
if os.environ.get("MKT_SKIP_KINASE_DICT"):
if bool_verbose:
logger.info(
"MKT_SKIP_KINASE_DICT set; skipping KinaseInfo deserialization."
)
if str_name is not None:
_deserialization_cache[str_name] = {}
return {}
if suffix not in DICT_FUNCS:
logger.error(
f"Serialization type ({suffix}) not supported; must be json, yaml, or toml."
)
return None
if str_path is None and suffix != "json":
logger.error("Only json deserialization is supported without providing a path.")
return None
if deserialization_kwargs is None:
deserialization_kwargs = DICT_FUNCS[suffix]["kwargs_deserialize"]
str_path = return_str_path_from_pkg_data(str_path)
dict_import = {}
if str_path.endswith(".tar.gz"):
dict_str = untar_files_in_memory(str_path, list_ids=list_ids)[1]
for val in tqdm(
dict_str.values(),
desc="Deserializing KinaseInfo objects in memory...",
bar_format=TQDM_BAR_FORMAT,
):
val_deserialized = DICT_FUNCS[suffix]["deserialize_str"](
val,
**deserialization_kwargs,
)
kinase_obj = kinase_schema.KinaseInfo.model_validate(val_deserialized)
dict_import[kinase_obj.hgnc_name] = kinase_obj
else:
list_file = glob.glob(os.path.join(str_path, f"*.{suffix}"))
for file in tqdm(
list_file,
desc="Deserializing KinaseInfo objects from files...",
bar_format=TQDM_BAR_FORMAT,
):
with open(file) as openfile:
val_deserialized = DICT_FUNCS[suffix]["deserialize_file"](
openfile,
**deserialization_kwargs,
)
kinase_obj = kinase_schema.KinaseInfo.model_validate(val_deserialized)
dict_import[kinase_obj.hgnc_name] = kinase_obj
if bool_remove:
clean_files_and_delete_directory(list_file)
dict_import = {key: dict_import[key] for key in sorted(dict_import.keys())}
if str_name is not None:
_deserialization_cache[str_name] = dict_import
return dict_import
STR_CONSERVATION_FILENAME = "KLIFSConservationData"
"""str: Basename (no suffix) of the persisted KLIFS conservation-data artifact."""
_conservation_cache = {}
"""dict: Module-level cache of loaded :class:`KLIFSConservationData` objects keyed by
``str_name`` (mirrors :data:`_deserialization_cache`)."""
[docs]
def serialize_conservation_data(
conservation_data: BaseModel,
suffix: str = "json",
str_path: str | None = None,
) -> str | None:
"""Serialize a :class:`KLIFSConservationData` artifact to a single file.
Unlike :func:`serialize_kinase_dict` (one file per kinase), this writes the whole
single-object artifact to ``{str_path}/KLIFSConservationData.{suffix}``, reusing the
:data:`DICT_FUNCS` serialization registry.
Parameters
----------
conservation_data : BaseModel
The :class:`mkt.schema.conservation_schema.KLIFSConservationData` object.
suffix : str
Serialization type supported: json, yaml, toml.
str_path : str | None
Directory to write into, by default None (the ``mkt.schema`` package directory).
Returns
-------
str | None
Path written, or None if the suffix is unsupported.
"""
if suffix not in DICT_FUNCS:
logger.error(
f"Serialization type ({suffix}) not supported; must be json, yaml, or toml."
)
return None
if os.name == "nt" and suffix == "toml":
logger.info("TOML serialization is not supported on Windows.")
return None
if str_path is None:
str_path = str(resources.files("mkt.schema"))
os.makedirs(str_path, exist_ok=True)
filepath = os.path.join(str_path, f"{STR_CONSERVATION_FILENAME}.{suffix}")
with open(filepath, "w") as outfile:
outfile.write(
DICT_FUNCS[suffix]["serialize"](
conservation_data.model_dump(),
**DICT_FUNCS[suffix]["kwargs_serialize"],
)
)
logger.info(f"Serialized KLIFSConservationData to {filepath}")
return filepath
[docs]
def load_conservation_data(
suffix: str = "json",
str_path: str | None = None,
str_name: str | None = None,
):
"""Load a :class:`KLIFSConservationData` artifact from a single file.
Parameters
----------
suffix : str
Deserialization type supported: json, yaml, toml.
str_path : str | None
Path to the artifact file, by default None (the packaged
``mkt/schema/KLIFSConservationData.json``).
str_name : str | None
If provided, cache the loaded object under this name in
:data:`_conservation_cache` to prevent reloading.
Returns
-------
KLIFSConservationData
The deserialized conservation-data object.
"""
if str_name is not None and str_name in _conservation_cache:
logger.info(f"Loading KLIFSConservationData from variable {str_name}...")
return _conservation_cache[str_name]
if suffix not in DICT_FUNCS:
logger.error(
f"Serialization type ({suffix}) not supported; must be json, yaml, or toml."
)
return None
from mkt.schema.conservation_schema import KLIFSConservationData
if str_path is None:
str_path = os.path.join(
str(resources.files("mkt.schema")), f"{STR_CONSERVATION_FILENAME}.{suffix}"
)
with open(str_path) as openfile:
val_deserialized = DICT_FUNCS[suffix]["deserialize_file"](
openfile,
**DICT_FUNCS[suffix]["kwargs_deserialize"],
)
obj = KLIFSConservationData.model_validate(val_deserialized)
if str_name is not None:
_conservation_cache[str_name] = obj
return obj
[docs]
def save_plot(
fig,
output_filename: str,
plot_type: str = "Plot",
bool_force_local: bool = True,
bool_image_subdir=True,
output_path: str | None = None,
bool_svg: bool = True,
bool_png: bool = True,
bool_pdf: bool = False,
**kwargs,
) -> None:
"""Save the current matplotlib figure in the requested vector/raster formats.
Parameters:
-----------
fig : matplotlib.figure.Figure
The figure object to save.
output_filename : str
Name of the output file to save the plot. Any extension is stripped; the
format suffixes are appended per the ``bool_svg``/``bool_png``/``bool_pdf`` flags.
plot_type : str
Description of the plot type for logging purposes (e.g., "Dynamic range plot")
bool_force_local : bool
If True, forces saving to the local dir (repo root or cwd) regardless of env var; default is True.
bool_image_subdir : bool
If True, saves images to a subdirectory named "images" within the output path; default is True.
output_path : str | None
Optional path to save the plot. If None, saves to the current working directory.
bool_svg : bool
If True, write an SVG copy; default is True.
bool_png : bool
If True, write a PNG copy; default is True.
bool_pdf : bool
If True, write a PDF copy; default is False.
**kwargs
Additional keyword arguments to pass to plt.savefig (e.g., {"dpi": 300}). Default is empty dict.
"""
import matplotlib.pyplot as plt
# remove extension if provided
output_filename = os.path.splitext(output_filename)[0]
# get_repo_root() > output_path > get_output_dir()
if bool_force_local:
if output_path is not None:
logger.info(
"bool_force_local is True, so ignoring provided output_path "
f"{output_path} and saving to local directory instead."
)
output_path = get_repo_root()
elif output_path is None:
output_path = get_output_dir()
# set default savefig parameters
savefig_params = {"bbox_inches": "tight"}
# update with any user-provided kwargs
savefig_params.update(kwargs)
suffixes = [
suffix
for suffix, flag in (("svg", bool_svg), ("png", bool_png), ("pdf", bool_pdf))
if flag
]
for suffix in suffixes:
if bool_image_subdir:
file_path = os.path.join(
output_path, "images", f"{output_filename}.{suffix}"
)
else:
file_path = os.path.join(output_path, f"{output_filename}.{suffix}")
fig.savefig(file_path, format=suffix, **savefig_params)
logger.info(f"{plot_type} saved to {file_path}")
plt.close(fig)