Source code for abl

"""
Python package for managing case files.


**Sub-packages**

.. autosummary::
   :toctree:

   templates

"""
import importlib.resources
import os
import pkgutil
import shutil
import sys
from pathlib import Path
from socket import gethostname

from abl.templates import makefile_usr
from eturb import logger, mpi
from fluidsim.base.output.base import OutputBase

# Sources for inclusion to makefile_usr.inc
# Dict[directory]  -> list of source files
makefile_usr_sources = {
    "toolbox": [
        ("frame.f", "FRAMELP"),
        ("mntrlog_block.f", "MNTRLOGD"),
        ("mntrlog.f", "MNTRLOGD"),
        ("mntrtmr_block.f", "MNTRLOGD", "MNTRTMRD"),
        ("mntrtmr.f", "MNTRLOGD", "MNTRTMRD", "FRAMELP"),
        ("rprm_block.f", "RPRMD"),
        ("rprm.f", "RPRMD", "FRAMELP"),
        ("io_tools_block.f", "IOTOOLD"),
        ("io_tools.f", "IOTOOLD"),
        ("chkpoint.f", "CHKPOINTD"),
        ("chkpt_mstp.f", "CHKPTMSTPD", "CHKPOINTD"),
        ("map2D.f", "MAP2D", "FRAMELP"),
        ("stat.f", "STATD", "MAP2D", "FRAMELP"),
        ("stat_IO.f", "STATD", "MAP2D", "FRAMELP"),
        ("math_tools.f",),
    ]
}

# Object files to be included in compilation
# Should be exported as USR environment variable
makefile_usr_obj = []
for objs in (
    (sources[0].replace(".f", ".o") for sources in list_of_sources)
    for list_of_sources in makefile_usr_sources.values()
):
    makefile_usr_obj.extend(objs)


[docs]def get_configfile(): """Get path of the Snakemake configuration file for the current machine. All configuration files are stored under ``etc`` sub-package. """ host = os.getenv("SNIC_RESOURCE", os.getenv("GITHUB_WORKFLOW", gethostname())) root = get_root() return root / "etc" / f"{host}.yml"
[docs]def get_root(): """Get the path to the current package, ``abl``.""" # Better than # >>> root = Path(__file__).parent? with importlib.resources.path(__name__, "__init__.py") as f: return f.parent
[docs]class Output(OutputBase): """Container and methods for getting paths of and copying case files. A path object :code:`self.root` points to the directory containing the case files. """ @staticmethod def _complete_info_solver(info_solver): """Complete the ParamContainer info_solver.""" OutputBase._complete_info_solver(info_solver) classes = info_solver.classes.Output.classes classes.PrintStdOut.module_name = "eturb.output.print_stdout" classes.PrintStdOut.class_name = "PrintStdOut" classes.PhysFields.module_name = "eturb.output.phys_fields" classes.PhysFields.class_name = "PhysFields" @staticmethod def _complete_params_with_default(params, info_solver): """This static method is used to complete the *params* container. """ # Bare minimum attribs = { "ONLINE_PLOT_OK": True, "period_refresh_plots": 1, "HAS_TO_SAVE": True, "sub_directory": "", } params._set_child("output", attribs=attribs) def __init__(self, sim=None): self.sim = sim try: self.name_solver = sim.info.solver.short_name except AttributeError: pass # Same as package name __name__ self.name_pkg = self.__module__ self.root = get_root() self._blacklist = {"prefix": "__", "suffix": (".vimrc", ".tar.gz", ".o")} if sim: super().__init__(sim) def _get_resources(self, name_pkg=None): """Get a generator of resources (files) in a package, excluding directories (subpackages). :returns: generator """ blacklist = self._blacklist if not name_pkg: name_pkg = self.name_pkg return ( f for f in importlib.resources.contents(name_pkg) if ( importlib.resources.is_resource(name_pkg, f) and not any(f.startswith(ext) for ext in blacklist["prefix"]) and not any(f.endswith(ext) for ext in blacklist["suffix"]) ) ) def _get_subpackages(self): """Get a dictionary of subpackages with values generated by :func:`get_resources`. :returns: dict """ path_pkg = self.root name_pkg = self.name_pkg return { subpkg.name.lstrip(f"{self.root.name}."): self._get_resources(subpkg.name) for subpkg in pkgutil.walk_packages([str(path_pkg)], prefix=f"{name_pkg}.") }
[docs] def get_paths(self): """Get a list of paths to all case files. :returns: list """ paths = [] # abl.usr -> /path/to/abl/abl.usr paths += [self.root / resource for resource in self._get_resources()] for subpkg, resources in self._get_subpackages().items(): # toolbox -> /path/to/abl/toolbox subpkg_root = self.root / subpkg.replace(".", os.sep) # main.f -> /path/to/abl/toolbox/main.f paths += [subpkg_root / resource for resource in resources] return paths
[docs] def copy(self, new_dir, force=False): """Copy case files to a new directory. The directory does not have to be present. :param new_dir: A str or Path-like instance pointing to the new directory. :param force: Force copy would overwrite if files already exist. """ # Avoid race conditions! Should be only executed by rank 0. if mpi.rank != 0: return abs_paths = self.get_paths() subpackages = self._get_subpackages() root = self.root def conditional_ignore(src, names): """Ignore if not found in ``abs_paths``.""" src = Path(src) include = abs_paths + [root / subpkg for subpkg in subpackages] exclude = tuple( name for name in names if not any((src / name) == path for path in include) ) logger.debug( "".join( ( f"- src: {src}", "\n- include:\n", "\n ".join(str(name) for name in include), "\n- exclude:\n", "\n ".join(exclude), "\n----", ) ) ) return exclude new_root = Path(new_dir) try: logger.info("Copying with shutil.copytree ...") copytree_kwargs = dict( src=root, dst=new_root, symlinks=True, ignore=conditional_ignore ) # Python 3.8+ shutil.copytree(**copytree_kwargs, dirs_exist_ok=True) except TypeError: try: logger.warning( "Python < 3.8: shutil.copytree may not proceed if directories exist." ) # Hoping that new_root has not been created shutil.copytree(**copytree_kwargs) except FileExistsError as e: logger.warning(e) logger.info("Copying with shutil.copy2 ...") # Copy one by one from the scratch if not new_root.exists(): logger.debug(f"Creating {new_root} ...") os.makedirs(new_root, exist_ok=True) for abs_path in abs_paths: rel_path = abs_path.relative_to(root) new_path = new_root / rel_path if not new_path.parent.exists(): os.makedirs(new_path.parent) logger.debug(f"Copying {new_path}") if new_path.exists(): if force: logger.warning(f"{new_path} would be overwritten ...") else: logger.warning( f"{new_path} exists, skipping. Use force=True to overwrite." ) continue shutil.copy2(abs_path, new_path) finally: logger.info(f"Copied: {root} -> {new_root}")
[docs] def write_makefile_usr(self, fp=sys.stdout, comments=""): """Write the makefile_usr.inc file which is required for compilation.""" comments += "\nAutogenerated using abl.Output.write_makefile_usr()" for path_dir, list_of_sources in makefile_usr_sources.items(): list_of_relative_sources = [] for sources in list_of_sources: list_of_relative_sources.append( [f"{path_dir}/{file}" for file in sources] ) output = makefile_usr.render( list_of_sources=list_of_relative_sources, comments=comments ) fp.write(output)