# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Utilities to create scripts and command-line tools."""
import ast
import codecs
import operator
import os.path
import functools
import types
import warnings
import numpy as np
from base64 import urlsafe_b64encode
from pathlib import Path
from uuid import uuid4
import yaml
from gammapy.utils.check import add_checksum, verify_checksum
__all__ = [
"from_yaml",
"get_images_paths",
"make_path",
"read_yaml",
"recursive_merge_dicts",
"to_yaml",
"write_yaml",
]
PATH_DOCS = Path(__file__).resolve().parent / ".." / ".." / "docs"
SKIP = ["_static", "_build", "_checkpoints", "docs/user-guide/model-gallery/"]
YAML_FORMAT = dict(sort_keys=False, indent=4, width=80, default_flow_style=False)
[docs]
def get_images_paths(folder=PATH_DOCS):
"""Generator yields a Path for each image used in notebook.
Parameters
----------
folder : str
Folder where to search.
"""
for i in Path(folder).rglob("images/*"):
if not any(s in str(i) for s in SKIP):
yield i.resolve()
[docs]
def from_yaml(text, sort_keys=False, checksum=False):
"""Read YAML file.
Parameters
----------
text : str
yaml str
sort_keys : bool, optional
Whether to sort keys. Default is False.
checksum : bool
Whether to perform checksum verification. Default is False.
Returns
-------
data : dict
YAML file content as a dictionary.
"""
data = yaml.safe_load(text)
checksum_str = data.pop("checksum", None)
if checksum:
yaml_format = YAML_FORMAT.copy()
yaml_format["sort_keys"] = sort_keys
yaml_str = yaml.dump(data, **yaml_format)
if not verify_checksum(yaml_str, checksum_str):
warnings.warn("Checksum verification failed.", UserWarning)
return data
[docs]
def read_yaml(filename, logger=None, checksum=False):
"""Read YAML file.
Parameters
----------
filename : `~pathlib.Path`
Filename.
logger : `~logging.Logger`
Logger.
checksum : bool
Whether to perform checksum verification. Default is False.
Returns
-------
data : dict
YAML file content as a dictionary.
"""
path = make_path(filename)
if logger is not None:
logger.info(f"Reading {path}")
text = path.read_text()
return from_yaml(text, checksum=checksum)
[docs]
def to_yaml(dictionary, sort_keys=False):
"""Dictionary to yaml file.
Parameters
----------
dictionary : dict
Python dictionary.
sort_keys : bool, optional
Whether to sort keys. Default is False.
"""
from gammapy.utils.metadata import CreatorMetaData
yaml_format = YAML_FORMAT.copy()
yaml_format["sort_keys"] = sort_keys
text = yaml.safe_dump(dictionary, **yaml_format)
creation = CreatorMetaData()
return text + creation.to_yaml()
[docs]
def write_yaml(
text, filename, logger=None, sort_keys=False, checksum=False, overwrite=False
):
"""Write YAML file.
Parameters
----------
text : str
yaml str
filename : `~pathlib.Path`
Filename.
logger : `~logging.Logger`, optional
Logger. Default is None.
sort_keys : bool, optional
Whether to sort keys. Default is True.
checksum : bool, optional
Whether to add checksum keyword. Default is False.
overwrite : bool, optional
Overwrite existing file. Default is False.
"""
if checksum:
text = add_checksum(text, sort_keys=sort_keys)
path = make_path(filename)
path.parent.mkdir(exist_ok=True)
if path.exists() and not overwrite:
raise IOError(f"File exists already: {path}")
if logger is not None:
logger.info(f"Writing {path}")
path.write_text(text)
def make_name(name=None):
"""Make a dataset name."""
if name is None:
name = urlsafe_b64encode(codecs.decode(uuid4().hex, "hex")).decode()[:8]
while name[0] == "_":
name = urlsafe_b64encode(codecs.decode(uuid4().hex, "hex")).decode()[:8]
if not isinstance(name, str):
raise ValueError(
"Name argument must be a string, "
f"got '{name}', which is of type '{type(name)}'"
)
return name
[docs]
def make_path(path):
"""Expand environment variables on `~pathlib.Path` construction.
Parameters
----------
path : str, `pathlib.Path`
Path to expand.
"""
# TODO: raise error or warning if environment variables that don't resolve are used
# e.g. "spam/$DAMN/ham" where `$DAMN` is not defined
# Otherwise this can result in cryptic errors later on
if path is None:
return None
else:
return Path(os.path.expandvars(path))
[docs]
def recursive_merge_dicts(a, b):
"""Recursively merge two dictionaries.
Entries in 'b' override entries in 'a'. The built-in update function cannot be
used for hierarchical dicts, see:
http://stackoverflow.com/questions/3232943/update-value-of-a-nested-dictionary-of-varying-depth/3233356#3233356
Parameters
----------
a : dict
Dictionary to be merged.
b : dict
Dictionary to be merged.
Returns
-------
c : dict
Merged dictionary.
Examples
--------
>>> from gammapy.utils.scripts import recursive_merge_dicts
>>> a = dict(a=42, b=dict(c=43, e=44))
>>> b = dict(d=99, b=dict(c=50, g=98))
>>> c = recursive_merge_dicts(a, b)
>>> print(c)
{'a': 42, 'b': {'c': 50, 'e': 44, 'g': 98}, 'd': 99}
"""
c = a.copy()
for k, v in b.items():
if k in c and isinstance(c[k], dict):
c[k] = recursive_merge_dicts(c[k], v)
else:
c[k] = v
return c
def requires_module(module_name):
"""
Decorator that conditionally enables a method or property based on the availability of a module.
If the specified module is available, the decorated method or property is returned as-is.
If the module is not available:
- For methods: replaces the method with one that raises ImportError when called.
- For properties: replaces the property with one that raises ImportError when accessed.
Parameters
----------
module_name : str
The name of the module to check for.
Returns
-------
function or property
The original object if the module is available, otherwise a fallback.
"""
def decorator(obj):
try:
__import__(module_name)
return obj # Module is available
except ImportError:
if isinstance(obj, property):
return property(
lambda self: raise_import_error(module_name, is_property=True)
)
elif isinstance(obj, (types.FunctionType, types.MethodType)):
@functools.wraps(obj)
def wrapper(*args, **kwargs):
raise_import_error(module_name)
return wrapper
else:
raise TypeError(
"requires_module can only be used on methods or properties."
)
return decorator
def raise_import_error(module_name, is_property=False):
"""
Raises an ImportError with a descriptive message about a missing module.
Parameters
----------
module_name : str
The name of the required module.
is_property : bool
Whether the error is for a property (affects the error message).
"""
kind = "property" if is_property else "method"
raise ImportError(f"The '{module_name}' module is required to use this {kind}.")
# Mapping of AST operators to NumPy functions
_OPERATORS = {
ast.And: np.logical_and,
ast.Or: np.logical_or,
ast.Eq: operator.eq,
ast.NotEq: operator.ne,
ast.Lt: operator.lt,
ast.LtE: operator.le,
ast.Gt: operator.gt,
ast.GtE: operator.ge,
ast.In: lambda a, b: np.isin(a, b),
ast.NotIn: lambda a, b: ~np.isin(a, b),
}
def logic_parser(table, expression):
"""
Parse and apply a logical expression to filter rows from an Astropy Table.
This function evaluates a logical expression on each row of the input table
and returns a new table containing only the rows that satisfy the expression.
The expression can reference any column in the table by name and supports
logical operators (`and`, `or`), comparison operators (`<`, `>`, `==`, `!=`, `in`),
lists, and constants.
Parameters
----------
table : `~astropy.table.Table`
The input table to filter.
expression : str
The logical expression to evaluate on each row. The expression can reference
any column in the table by name.
Returns
-------
table : `~astropy.table.Table`
A table view containing only the rows that satisfy the expression. If no rows
match the condition, an empty table with the same column names and data types
as the input table is returned.
Examples
--------
Given a table with columns 'OBS_ID' and 'EVENT_TYPE':
>>> from astropy.table import Table
>>> data = {'OBS_ID': [1, 2, 3, 4], 'EVENT_TYPE': ['1', '3', '4', '2']}
>>> table = Table(data)
>>> expression = '(OBS_ID < 3) and (OBS_ID > 1) and ((EVENT_TYPE in ["3", "4"]) or (EVENT_TYPE == "3"))'
>>> filtered_table = logic_parser(table, expression)
>>> print(filtered_table)
OBS_ID EVENT_TYPE
------ ----------
2 3
"""
def eval_node(node):
if isinstance(node, ast.BoolOp):
op_func = _OPERATORS[type(node.op)]
values = [eval_node(v) for v in node.values]
result = values[0]
for v in values[1:]:
result = op_func(result, v)
return result
elif isinstance(node, ast.Compare):
left = eval_node(node.left)
for op, comparator in zip(node.ops, node.comparators):
right = eval_node(comparator)
op_func = _OPERATORS[type(op)]
left = op_func(left, right)
return left
elif isinstance(node, ast.Name):
if node.id not in table.colnames:
raise KeyError(
f"Column '{node.id}' not found in the table. Available columns: {table.colnames}"
)
return table[node.id]
elif isinstance(node, ast.Constant):
return node.value
elif isinstance(node, ast.List):
return [eval_node(elt) for elt in node.elts]
else:
raise ValueError(f"Unsupported expression type: {type(node)}")
expr_ast = ast.parse(expression, mode="eval")
mask = eval_node(expr_ast.body)
return table[mask]
